[백준] 2606 - 바이러스 (py) + DFS 구현

문제

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

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

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

코드

import sys

com_num = int(input())
link_num = int(input())

dic = {}
for i in range(link_num):
    a, b = map(int, sys.stdin.readline().split())
    dic[a] = dic.get(a, []) + [b]
    dic[b] = dic.get(b, []) + [a]


def dfs(s, dict):
    stack_dfs= []
    visited = set()
    start_node = s
    stack_dfs.append(start_node)
    if s not in dic.keys():
        return [s]
    while stack_dfs:
        node = stack_dfs.pop()
        if node not in visited:
            visited.add(node)
            stack_dfs.extend(dict[node])
    return visited

print(len(dfs(1,dic))-1)

풀이

연결되어 있는 모든 컴퓨터의 수를 찾아야 하므로 DFS가 적절하다.
DFS을 구현 시 생각해야 할 것들은 다음과 같다.

1. 재귀적 구현

def dfs(dict, v, visited=[]):
    visited.append(v)
    for node in dict[v]:
        if node not in visited:
            dfs(dict, node, visited)
    return visited

노드들이 각각 key: value로 연결된 딕셔너리를 만들었다. visited 배열에 시작 노드를 추가해 주고, 그 노드와 연결된 다른 노드들에 대해 재귀적으로 dfs를 반복한다. 이때 모든 노드들이 방문되었다면 끝낸다. (return)

2. 반복문 구현

def dfs(s, dict):
    stack_dfs= []
    visited = set()
    start_node = s
    stack_dfs.append(start_node)
    if s not in dic.keys():
        return [s]
    while stack_dfs:
        node = stack_dfs.pop()
        if node not in visited:
            visited.add(node)
            stack_dfs.extend(dict[node])
    return visited

스택을 이용한다. 방문을 진행할 노드들을 담는 스택을 만들고, visted는 set으로 구현하였다. startnode를 정해주고, stack에 start node를 넣고 반복문을 시행한다. stack에 쌓인 노드들에 대해 pop을 시행해 주고, 만일 node가 visited에 없다면 visited에 추가한다. stack_dfs에 나머지 방문해야 할 노드들을 전부 넣어준다.