[week03] 백준 11725번 python 풀이

문제

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

11725번: 트리의 부모 찾기

왜 틀렸을까?

이 문제를 풀기 전에는 이차원 배열에서 0과 1을 통해 board[1][3] 이면 1에서 3을 갈 수 있다는 식으로 연결 상태를 나타냈었는데, 이 문제는 노드의 개수가 최대 100,000개 이므로 이차원 배열에서 연결상태를 표시하게 되면 10만*10만 = 100억개의 칸을 사용하기 때문에 메모리 초과가 날 수 밖에 없다.

# 원래 알던 이차원배열에 모두 저장하는 방식
board = [[0]*(n+1) for i in range(n+1)]

for i in range(n-1):
    a, b = map(int, sys.stdin.readline().split())
    board[a][b] = 1
    board[b][a] = 1

따라서 새로운 방식이 필요했는데, 그 방식은 같은 이차원 배열을 사용하지만 그래프의 연결 상태만 표시해준다면 메모리의 사용을 크게 줄일 수 있다.

# 연결 상태만 표시하는 새로운 방식!
graph = [[] for i in range(N+1)]
for _ in range(N-1):
    a, b = map(int, sys.stdin.readline().split())
    graph[a].append(b)
    graph[b].append(a)

메모리 초과를 해결한 다음에는 시간초과가 떴는데 그 이유는 기존 0,1을 사용한 이차원 배열에서 사용하던 방식을 새로운 연결 방식에 그대로 사용했기 때문이다.

# 시간초과 원인
visited = [0]*(n+1)
arr = []

def dfs(s, prev):
    arr.append((s, prev))
    visited[s] = 1
    for i in range(1, n+1):
        if visited[i] == 0 and i in graph[s]:
            dfs(i, s)

시간초과의 원인은 dfs가 돌면서 모든 노드들을 탐색하는데 그 노드를 탐색할 때마다 for 문을 10만번씩 돌리기 때문에 연산을 100억번 해야 하기 때문이다.

 

중요 포인트

최종적으로 이 문제는 DFS로 풀 수도 있고, BFS로 풀 수도 있다.

 

DFS 방식

문제에서 트리의 루트를 1이라고 정했기 때문에 1부터 DFS 탐색을 시작하고, 1과 연결되어 있는 Node들 중에 방문하지 않은 노드를 방문하는데, 이 때 visited배열의 index에 탐색을 시작한 node(즉, 부모 노드) 를 저장한다. 이렇게 되면 visited에 0이 아닌 수가 저장이 되기 때문에 다시 방문하지 않는다.

 

DFS를 사용한 풀이의 핵심은 DFS는 항상 부모에서 자식으로 이동한다는 점이다.

 

BFS 방식

1부터 시작하기 때문에 1을 queue에 넣고 while문을 돌린다. while문 안에서 빠져나오면 그 값을 now라고 하고, now와 연결되어있는 node들 중에서 가보지 않는 노드들을 가서 ans[nxt]의 값에 now를 넣으면 된다.

 

BFS를 이용한 방식도 DFS와 유사한 것을 알 수 있다.

코드 - DFS

import sys
sys.setrecursionlimit(10**6)
n = int(sys.stdin.readline())

graph = [[] for i in range(n+1)]

for i in range(n-1):
    a, b = map(int, sys.stdin.readline().split())
    graph[a].append(b)
    graph[b].append(a)

visited = [0]*(n+1)

arr = []

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

dfs(1)

for x in range(2, n+1):
    print(visited[x])

코드 - BFS

import sys
from collections import deque

N = int(sys.stdin.readline())
graph = [[] for i in range(N+1)]
for _ in range(N-1):
    a, b = map(int, sys.stdin.readline().split())
    graph[a].append(b)
    graph[b].append(a)

queue = deque()
queue.append(1)

ans = [0]*(N+1)

def bfs():
    while queue:
        now = queue.popleft()
        for nxt in graph[now]:
            if ans[nxt] == 0:
                ans[nxt] = now
                queue.append(nxt)

bfs()
res = ans[2:]
for x in res:
    print(x)