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

1. 문제

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

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

 

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

 

시작 노드가 1인 그래프가 주어졌을 때, 전체 노드 가운데 1과 연결된 노드를 모두 탐색하는 문제입니다.

이때, 1과 연결되어 있지 않은 노드(ex. 4, 7)가 존재할 수 있습니다.

---

2. 입력 및 출력

첫째 줄에는 컴퓨터의 수가 주어진다. 컴퓨터의 수는 100 이하인 양의 정수이고 각 컴퓨터에는 1번부터 차례대로 번호가 매겨진다. 둘째 줄에는 네트워크 상에서 직접 연결되어 있는 컴퓨터 쌍의 수가 주어진다. 이어서 그 수만큼 한 줄에 한 쌍씩 네트워크 상에서 직접 연결되어 있는 컴퓨터의 번호 쌍이 주어진다.

 

1번 컴퓨터가 웜 바이러스에 걸렸을 때, 1번 컴퓨터를 통해 웜 바이러스에 걸리게 되는 컴퓨터의 수를 첫째 줄에 출력한다.

 

노드(컴퓨터)의 총 개수와, 직접 연결된 노드 쌍의 수, 노드 쌍에 대한 정보가 차례로 입력됩니다.

입력된 정보를 토대로 1번 노드와 연결된 노드들의 수를 구해 아래와 같이 출력하면 됩니다.

 

> 예제 입력

7
6
1 2
2 3
1 5
5 2
5 6
4 7

 

1과 연결된 노드는 2, 3, 4, 5로 총 4대의 컴퓨터가 바이러스에 걸린다.

> 예제 출력

4

---

3. 제출 코드

그래프의 시작 노드(1)로부터 연결된 노드까지 탐색하는 문제인데, 이 경우 깊이 우선 탐색(DFS)과 너비 우선 탐색(BFS)을 사용해 문제를 해결할 수 있습니다.

 

틀렸습니다와 런타임 에러(KeyError)가 나왔는데, 1) 방향 설명이 없는 경우 양방향으로 edge를 설정해야 하고, 2) 파이썬 리스트가 아닌 dictionary를 사용하는 경우 Key error에 주의해야 합니다.

 

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

from collections import defaultdict

def dfs(graph, visited, i):
    visited.append(i)

    for node in graph[i]:
        if node not in visited:
            dfs(graph, visited, node)

    return len(visited) - 1


### Main ###
N = int(input())
M = int(input())

graph = defaultdict(list)
for i in range(M):
    key, val = map(int, input().split())
    graph[key].append(val)
    graph[val].append(key)

visited = []
print(dfs(graph, visited, 1))

 

오류 발생으로 인해 기존 코드들을 리뷰하던 중, defaultdict(list) 라는 함수를 발견했습니다.

그래프 딕셔너리를 만들던 이전 코드에서는 key가 그래프에 있는지 확인하고, 있다면 리스트에 노드를 추가하는 방식으로 코드를 작성해 가독성이 떨어졌습니다. (이전 문제에서는 맞았습니다가 나왔지만, 이번 문제에서는 숨겨진 테스트 케이스에서 key 처리를 잘하지 못한 것 같습니다.. KeyError가 발생했어요)

# 이전 코드
graph = dict()
for i in range(M):
    key, val = map(int, input().split())
    if key not in graph:
        graph[key] = [val]
    else:
        graph[key].append(val)

    if val not in graph:
        graph[val] = [key]
    else:
        graph[val].append(key)
        
# 수정 후 코드
graph = defaultdict(list)
for i in range(M):
    key, val = map(int, input().split())
    graph[key].append(val)
    graph[val].append(key)

 

DFS는 1) 스택을 사용하여 구현하거나, 2) recursive 하게 구현이 가능합니다.

 

스택을 python list를 사용해 구현한 코드는 아래 포스트에서 확인할 수 있습니다.

다른 코드에서는 list의 pop을 사용해 구현하기도 했는데, 내장 함수보다는 0번 요소에 직접 접근하는 것이 더 직관적이어서 아래 코드는 arr[0]에 접근하는 방식으로 짜여 있습니다.

2024.04.10 - [코딩테스트] - [백준]1260번 DFS와 BFS (Python/파이썬) - 런타임 에러(KeyError) 해결

[백준]1260번 DFS와 BFS (Python/파이썬) - 런타임 에러(KeyError) 해결

 

이번 제출 코드에서는 DFS를 recursive 하게 구현했습니다.

def dfs(graph, visited, i):
    visited.append(i)

    for node in graph[i]:
        if node not in visited:
            dfs(graph, visited, node)

    return len(visited) - 1

 

그래프에서 방문해야 할 노드 후보 리스트를 visited로 선언하고, current node i에 1(start node)을 넣어줍니다.

그래프 딕셔너리의 key i의 value 노드들 가운데, 아직 방문하지 않은(if node not in visited) 노드가 있다면, recursive call로 dfs 함수를 호출하여 방문합니다.

이런 식으로 노드를 탐색하면, 1을 입력했을 때, for문에서 2에 먼저 접근이 되고, dfs(graph, visited, 2)를 호출하여 3에 방문하게 됩니다. recursive call을 통해 1 → 2 → 3에 방문하고 나면, 다시 5로 돌아와 6을 방문하게 됩니다.

이렇게 하나의 branch를 다 탐색한 이후, 옆으로 넘어가는 방식을 DFS라고 부릅니다.

 

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

from collections import deque, defaultdict

def bfs(graph):
    need_visit = deque([])
    visited = []

    need_visit.append(1)
    while need_visit:
        node = need_visit.popleft()

        if node not in visited:
            visited.append(node)
            need_visit.extend(graph[node])

    return len(visited) - 1

### Main ###
N = int(input())
M = int(input())

graph = defaultdict(list)
for i in range(M):
    key, val = map(int, input().split())
    graph[key].append(val)
    graph[val].append(key)

print(bfs(graph))

 

너비 우선 탐색은 큐(Queue)를 사용해 구현합니다.

자세한 구현 방식은 이전 포스트에서도 다뤘습니다.

2024.04.11 - [코딩테스트] - [백준]2178번 미로탐색 (Python/파이썬) - BFS 코테 문제

[백준]2178번 미로탐색 (Python/파이썬) - BFS 코테 문제

 

 

3) 통합 코드

from collections import deque, defaultdict

def bfs(graph):
    need_visit = deque([])
    visited = []

    need_visit.append(1)
    while need_visit:
        node = need_visit.popleft()

        if node not in visited:
            visited.append(node)
            need_visit.extend(graph[node])

    return len(visited) - 1


def dfs(graph, visited, i):
    visited.append(i)

    for node in graph[i]:
        if node not in visited:
            dfs(graph, visited, node)

    return len(visited) - 1


### Main ###
N = int(input())
M = int(input())

graph = defaultdict(list)
for i in range(M):
    key, val = map(int, input().split())
    graph[key].append(val)
    graph[val].append(key)

print(bfs(graph))

#visited = []
#print(dfs(graph, visited, 1))

---

BFS, DFS 언제 이해하나 했는데 슬슬 감이 오는 거 같습니다

학교 다닐 때 열심히 했으면 더 좋았겠지만, 어쩔 수 없죠 뭐 ㅎㅎ 매일 하나씩 열심히 해보렵니다

---

문제 출처

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

2606번: 바이러스