[Python] 백준 11725 트리의 부모찾기

문제

루트 없는 트리가 주어진다. 이때, 트리의 루트를 1이라고 정했을 때, 각 노드의 부모를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 노드의 개수 N (2 ≤ N ≤ 100,000)이 주어진다. 둘째 줄부터 N-1개의 줄에 트리 상에서 연결된 두 정점이 주어진다.

출력

첫째 줄부터 N-1개의 줄에 각 노드의 부모 노드 번호를 2번 노드부터 순서대로 출력한다.

 

11725번: 트리의 부모 찾기

---

나의 풀이 (시간 초과)

왜 시간 초과가 났을까? While문 안에 Not in 이라는 연산자 역시 O(n)의 시간 복잡도를 가지고 있으니, O(n2)가 된다. 

import sys
import collections

n = int(sys.stdin.readline())
graph = collections.defaultdict(list)
for _ in range(n-1):
    s, e = map(int, sys.stdin.readline().split())
    graph[s].append(e)
    graph[e].append(s)

visited = []
ans = []
def DFS(graph, root):
    stack = [root] # stack = 가야할 위치를 저장
    while stack: 
        node = stack.pop() 
        if node not in visited: # 방문 이력이 없다면
            visited.append(node) # 방문 이력에 기록을 남기고
            need_to_visit = list(set(graph[node])-set(visited))
            stack.extend(need_to_visit) # 해당 노드와 연결된 노드를 (방문 이력 있는곳을 제외한) stack에 저장
            for nod in need_to_visit: # 
                ans.append([nod, node])
DFS(graph, 1)
ans.sort()
for a in ans:
    print(a[1])

 

리스트값을 찾는 조건문은 줄이는게 좋다. 

그래서 조건문에 확인 내용이 숫자라면, 리스트의 인덱스와 엮어서 짜야 계산시간을 줄일 수 있다. 

import sys
import collections
sys.setrecursionlimit(10000000000)

n = int(sys.stdin.readline())
graph = collections.defaultdict(list)
for _ in range(n-1):
    s, e = map(int, sys.stdin.readline().split())
    graph[s].append(e)
    graph[e].append(s)

parents = [0 for i in range(n+1)]
def DFS(graph, node):
    # root들이 연결하고 있는 자식 노드를 확인하고, 저장하고 각각의 자식노드를 또 확인하고 저장하고 재귀..
    for child in graph[node]:
        if parents[child] == 0:
            parents[child] = node
            DFS(graph, child)

DFS(graph, 1)
for i in range(2, n+1):
    print(parents[i])