[Python] 2606 - 바이러스

1. 문제 정의

주어진 컴퓨터들의 연결상태에 대해, 1번 컴퓨터를 통해 바이러스에 감염되는 컴퓨터의 수를 구하는 문제이다.

그림과 같은 예제에 대해서, 1번 컴퓨터와 연결된 컴퓨터의 수는 2, 3, 5, 6 으로 총 4대이다.

 

2. 풀이 방식

해당 문제를 해결하기 위해 그래프의 탐색 알고리즘을 활용하였다. 탐색 알고리즘에는 두 가지 방법으로, DFS(깊이 우선 탐색)와 BFS(넓이 우선 탐색)이 있다.

 

1) 리스트 생성

# 2606 바이러스

n = int(input()) # 노드 갯수(컴퓨터)
v = int(input()) # 엣지 갯수(연결 선)

graph = [[] for i in range(n+1)] # 그래프 초기화
visited = [0] * (n+1) # 방문한 컴퓨터인지 표시

for i in range(v):
    a, b = map(int, input().split())
    graph[a] += [b] # a에 b 연결
    graph[b] += [a] # b에 a 연결 -> 양방향

graph는 n+1개의 리스트로 생성, visited는 n+1길이의 0으로 이루어진 리스트 생성한다.

연결된 갯수(v)만큼 반복하여 연결 관계를 graph에 저장한다.

graph에 저장된 리스트는 다음과 같다.

(0번 인덱스는 0번 컴퓨터)

1번 인덱스에 해당하는 숫자는 1번 컴퓨터와 연결된 컴퓨터를 의미, 2번 컴퓨터와 5번 컴퓨터와 연결되어 있음을 의미한다.

 

2) DFS

def dfs(num):
    visited[num] = 1
    
    for nx in graph[num]:
        if (visited[nx] == 0):
            dfs(nx)
dfs(1)
print(sum(visited) - 1)

방문할 컴퓨터 번호를 num로 입력 받은 후, visited에 방문 표시.

graph[num]은 위에서 알 수 있듯 num번 컴퓨터에 연결된 컴퓨터들이다. visited[nx] == 0을 통해 방문한 컴퓨터인지 검사 하고 방문하지 않은 컴퓨터인 경우 dfs(nx)를 통해 함수를 재귀 호출하여 이 과정을 반복한다.

이후 sum(visited)를 이용하여 방문한 컴퓨터들의 갯수를 합으로 구한 후 -1을 한다. (1번 컴퓨터를 제외한 연결된 컴퓨터 수)

 

3) BFS

from collections import deque

visited[1] = 1 # 1번부터 시작이기에 1번 표시
Q = deque([1])
while Q:
    c = Q.popleft()
    
    for nx in graph[c]:
        if (visited[nx] == 0):
            Q.append(nx)
            visited[nx] = 1

print(sum(visited) - 1)

BFS는 deque를 이용한다. visited[1] = 1로 1번 컴퓨터에 방문 표시한다.

Q = deque([1])로 1번 컴퓨터의 덱을 만든다. Q에 값이 있을 동안 while문이 동작하며 c = Q.popleft()를 통해 Q의 왼쪽값을 c로 받으면서 반복을 진행한다.

c번 컴퓨터와 연결된 컴퓨터에 대해 방문을 하지 않은 컴퓨터가 있을 경우, Q에 추가하고 방문 표시를 한다.

while문을 빠져나온 후 sum(visited)를 통해 갯수를 구한 후 -1을 하여 1번 컴퓨터를 빼준다.

 

3. 후기

처음으로 풀어본 그래프 탐색 문제였다. 지금까지는 단순한 구현, 정렬 등의 알고리즘만 해결하였으나 그래프 문제를 머리로만 해결해보았지 코드로 작성해보기는 처음이었다. 힘겹게 풀이를 찾아보며 문제를 해결하였다. 공부가 더 필요하다 생각한다.