[Baekjoon] 2606 - 바이러스 (Python)

문제

[BaekJoon] 2606 바이러스

컴퓨터 1대가 웜 바이러스에 걸리게되면 서로 연결되어 있는 컴퓨터도 같이 웜 바이러스에 걸리게된다.

주어진 컴퓨터에 대해 1번 컴퓨터와 연결된 컴퓨터가 총 몇 개인지 출력하는 문제이다.

Input

총 컴퓨터의 대수 N

컴퓨터 간 연결된 선의 개수 M

직접 연결되어 있는 M개 컴퓨터의 번호 쌍

Output

1번 컴퓨터가 웜 바이러스에 걸렸을 때, 1번 컴퓨터를 통해 웜 바이러스에 걸리게 되는 컴퓨터의 수

풀이

graph를 초기화 한다. n+1을 해주는 이유는 1번 인덱스를 1번 컴퓨터로 사용하기 위함이다. 즉 0번 인덱스를 비워두는 것이다. 방문했는지 확인하기 위한 visited도 마찬가지로 n+1개를 생성한다.

그 후에 입력한 직접 연결되어 있는 컴퓨터 번호 쌍들을 양방향으로 연결해주어 graph를 만든다.

방문하지 않은 컴퓨터는 dfs를 해주면서 방문한 컴퓨터는 1로 바꿔준다.

그리고 마지막에 1번 컴퓨터를 제외한 1번 컴퓨터와 연결된 개수를 출력하기 위해 -1을 해주고 출력하면 정답이다.

Code

import sys
input = sys.stdin.readline

n = int(input()) # 노드 개수(컴퓨터 수)
m = int(input()) # 연결선 수(간선 수)
graph = [[] for i in range(n+1)]
visited = [0] * (n+1)
for _ in range(m):
    a, b = map(int, input().strip().split())
    graph[a] += [b]
    graph[b] += [a]

def dfs(v):
    visited[v] = 1
    for i in graph[v]:
        if visited[i] == 0:
            dfs(i)

dfs(1)
print(sum(visited)-1) # 1번 컴퓨터를 제외하고 1번 컴퓨터와 연결된 개수 출력

    graph[a] += [b]
    graph[b] += [a]

이 부분을 알기 쉽게 Input에 따라 단계별로 정리하자면

1 2

[[], [2], [1], [], [], [], [], []]

2 3

[[], [2], [1, 3], [2], [], [], [], []]

1 5

[[], [2, 5], [1, 3], [2], [], [1], [], []]

5 2

[[], [2, 5], [1, 3, 5], [2], [], [1, 2], [], []]

5 6

[[], [2, 5], [1, 3, 5], [2], [], [1, 2, 6], [5], []]

4 7

[[], [2, 5], [1, 3, 5], [2], [7], [1, 2, 6], [5], [4]]