문제링크

BFS 풀이(메모리 초과)

루트 노드로부터 전체 노드를 탐색하며 다음 노드로 넘어갈 때 현재 노드의 값을 부모 노드로 기록하려고 했다.

그러나 그렇게 하면 메모리 초과가 발생한다. 왜냐하면, BFS는 큐를 사용해야 하기 때문이다.

함수

def bfs(graph,root):
    q = deque([root])
    visited = [False]*(N+1)
    root_list = [0]*(N+1)

    while q:
        cur = q.popleft()
        for i in range(1,N+1):
            if graph[cur][i] and not visited[i]:
                q.append(i)
                visited[i] = True
                root_list[i] = cur

    return root_list

root_list에는 부모 노드의 번호를 담았다.

예를 들어, root_list = [1,2,3,4,5]일 때 2번 노드의 부모노드는 3번 노드이다.

오답 코드

import sys
from collections import deque
input = sys.stdin.readline

def init():
    N = int(input())
    graph = [[0]*(N+1) for _ in range(N+1)]

    for _ in range(N-1):
        a,b = map(int,input().split())
        graph[a][b] = graph[b][a] = 1
    return N,graph

def bfs(graph,root):
    q = deque([root])
    visited = [False]*(N+1)
    root_list = [0]*(N+1)

    while q:
        cur = q.popleft()
        for i in range(1,N+1):
            if graph[cur][i] and not visited[i]:
                q.append(i)
                visited[i] = True
                root_list[i] = cur

    return root_list

if __name__ == '__main__':
    N, graph = init()
    root_list = bfs(graph,1)
    for x in root_list:
        print(x)

DFS 풀이(정답)

어떻게 하면 메모리를 줄일 수 있을까 고민하다가 재귀를 떠올렸다.

find_parent(재귀)

def find_parent(graph,root):
    cur = root
    for next in graph[cur]:
        if not parent_list[next]:
            parent_list[next] = cur
            find_parent(graph,next)
    return

parent_list는 앞서 소개한 root_list와 같은 역할이다.

코드를 짜다보니 변수명을 바꿔버림 🤣

아직 부모 노드가 반영되지 않은 노드들을 찾아서 부모 노드를 업데이트 해준다.

결론

메모리를 효율적으로 사용하고 싶을 때, DFS를 가져다 쓴다. (BFS 보다는!!)

코드

import sys
from collections import defaultdict
sys.setrecursionlimit(int(1e6))
input = sys.stdin.readline

def init():
    N = int(input())
    graph = defaultdict(list)

    for _ in range(N-1):
        a,b = map(int,input().split())
        graph[a].append(b)
        graph[b].append(a)
    return N,graph

def find_parent(graph,root):
    cur = root
    for next in graph[cur]:
        if not parent_list[next]:
            parent_list[next] = cur
            find_parent(graph,next)
    return

if __name__ == '__main__':
    N, graph = init()
    parent_list = [0]*(N+1)
    parent_list[1] = 1
    find_parent(graph,1)
    for x in parent_list[2:]:
        print(x)