[백준] 15681 트리와 쿼리, python

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

15681번: 트리와 쿼리

특정 노드를 정점으로 하는 서브 트리의 정점의 개수를 구해야한다.

무방향 그래프를 받아서 연결시켜주고, 루트 노드에서 시작해서 후위 순회로 자식 노드의 개수를 더해주면 된다.

이때 visit 로 방문했는지도 체크하고 해당 노드에서 자식 노드의 개수를 저장한다.

순회가 끝나면 visit에 자식 노드의 개수가 저장되어있으므로 특정 노드를 받아서 서브트리 내 정점 개수를 반환하면된다.

import sys; input = sys.stdin.readline
sys.setrecursionlimit = 10**9
linked_list = None
visit = None
# son_list = None
def postorder(cur_node):
    global linked_list, visit
    visit[cur_node] = 1
    num_of_son = 0
    for son in linked_list[cur_node]:
        if visit[son]: continue
        num_of_son += postorder(son)
    visit[cur_node] += num_of_son
    return visit[cur_node] # son + cur

def main():
    global linked_list, visit,son_list
    N, R, Q = map(int, input().split())
    linked_list = [[] for _ in range(N + 1)]
    visit = [0 for _ in range(N + 1)]
    # son_list = [0 for _ in range(N + 1)]
    for _ in range(N - 1):
        U, V = map(int, input().split())
        linked_list[U].append(V)
        linked_list[V].append(U)

    postorder(R)
    # print(visit)
    ans = ''
    for _ in range(Q):
        ans += str(visit[int(input())]) + '\n'
    print(ans)

if __name__ == "__main__":
    main()

dfs 여서 오랜만에 리커전에러를 봤다.

그리고 python은 통과했는데 pypy에서 메모리 초과가 났다! 흑흑...

C++은 밤이니까 안해야지...