[백준] 2606번 바이러스, 파이썬 문제 풀이, DFS와 BFS 차이

문제

백준 2606번 바이러스 문제를 해결했습니다. 1번 컴퓨터와 연결되어 있는 컴퓨터, 간접적으로 연결되어 있는 컴퓨터도 웜 바이러스에 감염됩니다. 이 때, 1번으로부터 감염된 컴퓨터가 몇 대인지 찾는 문제입니다.

https://www.acmicpc.net/problem/2606

2606번: 바이러스

 

문제 풀이

컴퓨터간 상하관계나 방향이 없기 때문에 그래프 문제로 해결할 수 있었습니다. 1번과 연결되어 있는 컴퓨터를 탐색하는 문제이므로 DFS와 BFS로 해결할 수 있습니다. 이번 문제는 DFS, BFS를 둘 다 적용해서 차이점을 확인해보도록 하겠습니다.

 

먼저, 2차원 배열에 연결된 컴퓨터의 정보를 저장합니다. 이는 그래프를 생성하는 과정입니다. 컴퓨터에 매겨져 있는 번호를 인덱스로 삼아, 연결되어 있는 컴퓨터의 정보를 저장합니다. 들어오는 한 쌍의 정보를 교차로 저장하면 연결된 컴퓨터의 정보를 저장할 수 있습니다.

 

그래프를 만들었으니, 탐색을 해야합니다. DFS, BFS 모두 방문 여부를 확인하고, 재방문을 하지 않는다는 특징이 있습니다. 따라서, 리스트를 만들어 방문 여부를 저장합니다. 

DFS (Depth First Search, 깊이 우선 탐색) 

DFS는 재귀 호출을 통해 구현할 수 있습니다. DFS 함수의 매개변수로 바이러스가 있는 컴퓨터의 번호를 받습니다. 위에서 만든 그래프를 보고, 해당 컴퓨터와 연결되어 있는 컴퓨터에 방문하지 않았다면 컴퓨터의 번호를 매개변수로 전달하면서 재귀호출을 합니다. 함수가 호출되는 수가 감염된 컴퓨터의 수입니다. 

DFS 함수는 다음과 같습니다.

cnt=0
def DFS(virus):
    global cnt
    visited[virus]=1

    for i in network[virus]:
        if (visited[i]==0):
            DFS(i)
            cnt+=1

BFS (Breadth First Search, 너비 우선 탐색)

BFS는 큐와 반복문을 통해 구현할 수 있습니다. 큐에 바이러스가 있는 컴퓨터를 넣고, 빼면서 이 컴퓨터와 연결되어 있는 컴퓨터에 방문하지 않았다면, 큐에 넣고 다시 반복문을 실행하는 것입니다. 마찬가지로, 큐에 들어가는 컴퓨터는 모두 감염이 된 컴퓨터이기 때문에 enqueue할 때마다 카운트를 해줍니다.

BFS 함수는 다음과 같습니다.

def BFS(virus):
    global cnt
    visited[virus] = 1
    queue = [virus]
    while queue:
        for i in network[queue.pop()]:
            if (visited[i]==0):
                visited[i]=1
                queue.insert(0, i)
                cnt+=1

전체 코드

전체적인 코드는 다음과 같습니다.

cnt=0
def DFS(virus):
    global cnt
    visited[virus]=1

    for i in network[virus]:
        if (visited[i]==0):
            DFS(i)
            cnt+=1

def BFS(virus):
    global cnt
    visited[virus] = 1
    queue = [virus]
    while queue:
        for i in network[queue.pop()]:
            if (visited[i]==0):
                visited[i]=1
                queue.insert(0, i)
                cnt+=1
################MAIN###############
N= int(input())
link = int(input())

network = [[]*(N+1) for _ in range(N+1)]

for i in range(link):
    a, b = map(int, input().split())
    network[a].append(b)
    network[b].append(a)

visited = [0]*(N+1)
BFS(1)
#DFS(1)
print(cnt)

결과

결과는 다음과 같이 나왔습니다.

▲ 각 코드의 결과

둘 다 메모리는 같지만, 시간은 BFS가 108ms, DFS가 104ms로 BFS가 살짝 더 느리게 나왔습니다. 하지만, 이론상 BFS가 조금 더 빠릅니다. 이와 같이 나온 이유는 제 코드가 최적화되어 있지 않기 때문이라고 생각합니다. 우선, 리스트로 작성이 되었기 때문에 중복된 값들이 들어갈 수 있습니다. 이는 튜플을 활용해서 시간을 더 줄일 수 있습니다. 다른 사람들의 코드를 참고해보면, 둘 다 실행 시간이 같은 것을 확인할 수 있었습니다. 

 

정리

DFS와 BFS는 상황에 맞게 사용해야 합니다. 대부분의 경우 DFS를 더 선호한다고 합니다. 저도 처음 구현할 때는 DFS부터 작성했던 것 같습니다. 더 많은 문제를 풀어보면서 조금씩 발전해 나가도록 하겠습니다.