[15681] 트리와 쿼리

 

> 깊이 우선 탐색

> 그래프 탐색

 

 

아래와 같이 코드를 작성해서

제출했더니

재귀 관련 오류가 났다. 

 

n, r, q = map(int, input().split())
tree = [[] for _ in range(n + 1)]
for i in range(n - 1):
    a, b = map(int, input().split())
    tree[a].append(b)
    tree[b].append(a)
query = [int(input()) for _ in range(q)]


def check(node):
    count[node] = 1
    for child in tree[node]:
        if not count[child]:
            check(child)
            count[node] += count[child]
    return count[node]


count = [0 for _ in range(n + 1)]
check(r)
for i in query:
    print(count[i])

 

setrecursionlimit 함수를 사용하여

재귀의 깊이를 늘려줬다.

10**9까지

그랬더니 메모리초과 오류가 떴다.

 

 

구글링을 해보니

재귀의 깊이를 너무 많이 늘려주면

메모리초과가 발생할 수 있다고 하여

아래와 같이 

변경해주었다.

 

input()도 sys.stdin.readline()으로 변경하여

시간 초과 문제도 해결하였다.

 

import sys
n, r, q = map(int, sys.stdin.readline().split())
tree = [[] for _ in range(n + 1)]
sys.setrecursionlimit(100000)

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

def check(node):
    count[node] = 1
    for child in tree[node]:
        if not count[child]:
            check(child)
            count[node] += count[child]
    return

count = [0] * (n+1)
check(r)
for _ in range(q):
    print(count[int(sys.stdin.readline())])