[파이썬/Python] 백준 2606번 바이러스

[파이썬/Python] 백준 2606번 바이러스

2606번: 바이러스

---

문제

신종 바이러스인 웜 바이러스는 네트워크를 통해 전파된다. 한 컴퓨터가 웜 바이러스에 걸리면 그 컴퓨터와 네트워크 상에서 연결되어 있는 모든 컴퓨터는 웜 바이러스에 걸리게 된다.

예를 들어 7대의 컴퓨터가 <그림 1>과 같이 네트워크 상에서 연결되어 있다고 하자. 1번 컴퓨터가 웜 바이러스에 걸리면 웜 바이러스는 2번과 5번 컴퓨터를 거쳐 3번과 6번 컴퓨터까지 전파되어 2, 3, 5, 6 네 대의 컴퓨터는 웜 바이러스에 걸리게 된다. 하지만 4번과 7번 컴퓨터는 1번 컴퓨터와 네트워크상에서 연결되어 있지 않기 때문에 영향을 받지 않는다.

어느 날 1번 컴퓨터가 웜 바이러스에 걸렸다. 컴퓨터의 수와 네트워크 상에서 서로 연결되어 있는 정보가 주어질 때, 1번 컴퓨터를 통해 웜 바이러스에 걸리게 되는 컴퓨터의 수를 출력하는 프로그램을 작성하시오.

 

해설

1. 간선의 양 끝점을 통해 정점끼리의 연결관계를 정리해준다.
2. 1번 컴퓨터가 숙주이므로, 1번 컴퓨터의 정점을 방문하고, 방문한 정점을 체크하여 연결관계를 확인한다. (DFS)
3. 더 이상 연결된 정점이 없으면 함수를 종료한다.
4. 1번 컴퓨터와 연결된 컴퓨터의 수를 출력해야하므로, 2번 컴퓨터부터 N번 컴퓨터까지 방문한 정점을 카운트해준다.

 

코드

import sys

N = int(sys.stdin.readline())
M = int(sys.stdin.readline())

network = {}
visited = [False for _ in range(N+1)]
count = 0

for _ in range(M):
    u, v = map(int, sys.stdin.readline().split())
    if u in network:
        network[u].append(v)
    else:
        network[u] = [v]
    
    if v in network:
        network[v].append(u)
    else:
        network[v] = [u]
        
def dfs(u):
    visited[u] = True
    if u in network:
        for v in network[u]:
            if visited[v] == False:
                dfs(v)

dfs(u = 1)

for u in range(2, N+1):
    if visited[u]:
        count += 1

print(count)

 

풀이

import sys

N = int(sys.stdin.readline())
M = int(sys.stdin.readline())

network = {}
visited = [False for _ in range(N+1)]
count = 0

1. 컴퓨터의 수, 네트워크 상에서 직접 연결되어 있는 컴퓨터 쌍의 수를 입력받는다.

네트워크 상에서 연결되어 있는 컴퓨터들을 딕셔너리를 통해 표현한다.

방문 표시를 위한 리스트와 감염된 컴퓨터의 수를 카운트할 변수를 초기화한다.

 

for _ in range(M):
    u, v = map(int, sys.stdin.readline().split())
    if u in network:
        network[u].append(v)
    else:
        network[u] = [v]
    
    if v in network:
        network[v].append(u)
    else:
        network[v] = [u]

2. 양 끝점 u, v를 입력받고, network 딕셔너리에 저장해준다.

양방향이기에 구분이 필요없다.

 

def dfs(u):
    visited[u] = True
    if u in network:
        for v in network[u]:
            if visited[v] == False:
                dfs(v)

3. 정점을 방문하는 함수 dfs 이다.

해당 정점을 방문했기 때문에 먼저 True 로 방문표시를 해준다.

정점 u 와 연결되어 있는 정점 v 들에 대해 dfs 함수를 재귀호출하여, 정점 u와 연결되어 있는 모든 정점들을 방문처리 해준다.

 

dfs(u = 1)

for u in range(2, N+1):
    if visited[u]:
        count += 1

print(count)

4. 1번 컴퓨터를 통해 감염될 컴퓨터들을 구해야하므로 u = 1 을 파라미터로 전달하여 dfs 함수를 호출한다.

1번 컴퓨터를 제외한 감염된 나머지 컴퓨터의 수를 구해야하므로 반복문을 통해 감염된 컴퓨터들(방문한 정점들)을 구한다.

 

피드백

바로 이전에 풀었던 문제와 매우 유사한 dfs 알고리즘 문제였기 때문에 쉽게 풀 수 있었다.

자세한 내용은 아래 글을 참고하자.

[파이썬/Python] 백준 11724번 연결 요소의 개수