[TIL 2024.07.19]Python DFS 활용 (2)

오늘 진행한 일들

• 알고리즘 강의 시청
• 알고리즘 코드 문제 풀이

알고리즘 문제

이번에도 지난 문제와 비슷한 DFS 활용 문제를 접했다.

바이러스(from 백준)

• 신종 바이러스인 웜 바이러스는 네트워크를 통해 전파된다.
• 한 컴퓨터가 웜 바이러스에 걸리면 그 컴퓨터와 네트워크 상에서 연결되어 있는 모든 컴퓨터는 웜 바이러스에 걸리게 된다.
• 1번 컴퓨터가 웜 바이러스에 걸렸을 때, 1번 컴퓨터를 통해 웜 바이러스에 걸리게 되는 컴퓨터의 수를 첫째 줄에 출력한다.

예시

작성한 코드

그래프 구성

graph = defaultdict(list)

for _ in range(link):
    x, y = map(int, sys.stdin.readline().split())
    graph[x].append(y)
    graph[y].append(x)

graph는 인접 리스트 형태로 그래프를 저장하고 각 컴퓨터 쌍을 입력받아 그래프에 추가하며 x와 y가 연결되어 있음을 나타내기 위해 양쪽 모두에 추가한다.

방문 배열 초기화 및 DFS 스택 초기화:

visited = defaultdict(bool)

dfs_stack = [1]
count = 0

visited는 각 컴퓨터가 방문되었는지 여부를 기록하는 배열이며, dfs_stack은 DFS 탐색을 위해 사용할 스택으로 초기값은 1번 컴퓨터이다. 또한 count는 감염된 컴퓨터의 수를 세기 위한 변수이다.

DFS 탐색:

while dfs_stack:
    top = dfs_stack.pop()
    visited[top] = True
    computer = graph[top]

    for next_node in computer:
        if not visited[next_node] and next_node not in dfs_stack:
            dfs_stack.append(next_node)
            count += 1

스택이 비어있지 않는 동안 반복하며, 스택에서 노드를 꺼내어 top에 저장한다. top이 방문되지 않았으면 방문을 기록하고 count를 증가시킨다. 또한 top의 인접 노드 중 방문되지 않은 노드를 스택에 추가한다.

전체 코드

import sys
from collections import defaultdict

n = int(sys.stdin.readline())
link = int(sys.stdin.readline())
graph = defaultdict(list)

for _ in range(link):
    x, y = map(int, sys.stdin.readline().split())
    graph[x].append(y)
    graph[y].append(x)

visited = defaultdict(bool)

dfs_stack = [1]
count = 0

while dfs_stack:
    top = dfs_stack.pop()
    visited[top] = True
    computer = graph[top]

    for next_node in computer:
        if not visited[next_node] and next_node not in dfs_stack:
            dfs_stack.append(next_node)
            count += 1
print(count)