[python] 11725번 - 트리의 부모 찾기

문제

파이썬 풀이

트리를 그래프를 저장하듯 리스트 형태로 저장하고, 부모 노드를 저장하는 리스트를 생성하는 것이 포인트였다.

부모 노드를 저장하는 리스트는 리스트의 인덱스가 해당 노드가 되고 저장된 값은 부모 노드가 된다.

 

예를 들어보면

2와 연결되어있는 노드 들은 2의 부모 노드와 자식 노드로 나뉜다.

이때 부모 노드를 저장하는 리스트에서 2의 부모 노드가 무엇인지 가져와 2와 인접한 노드들 중 부모 노드를 제외한다.

그럼 나머지는 전부 2의 자식노드가 된다. 이제 자식 노드로 구별된 노드들의 부모 노드를 2로 저장한다.

이와 같은 과정을 모든 노드들에 접근하면서 부모 노드를 저장한다.

모든 노드들을 접근하는 과정에서는 dfs나 bfs를 이용하여 전체 탐색한다.

 

나는 dfs를 통해 구현하였다.

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

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

parent = [0] * (n+1)

def dfs(v) : 
    for i in graph[v] :
        if parent[v] != i :
            parent[i] = v
            dfs(i)
            

dfs(1)
for i in range(2,n+1) :
    print(parent[i])

이렇게 제출했더니 런타임 에러 (RecursionError)가 나왔다.

그래서 열심히 찾아보던 중 재귀 깊이를 늘려주면 된다고 하여 윗 줄에  sys.setrecursionlimit(10 ** 9) 코드를 추가하였다.

from collections import deque
import sys
input = sys.stdin.readline
sys.setrecursionlimit(10 ** 9)

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

parent = [0] * (n+1)

def dfs(v) : 
    for i in graph[v] :
        if parent[v] != i :
            parent[i] = v
            dfs(i)
            

dfs(1)
for i in range(2,n+1) :
    print(parent[i])

이렇게 제출하니 이번에는 메모리 부족이 나왔다.

그런데 제출을 pypy3으로 했었는데 python3으로 하니 문제가 해결되었다.

왜인지는 잘 모르겠다. 이건 더 알아봐야겠다.