백준 15681 - 트리와 쿼리 (파이썬)

Updated:

우선 양방향 트리를 만든 다음에 dfs로 가장 말단의 노드까지 탐색했습니다.
그 노드부터 1씩 더해주면 size 배열에는 특정 정점을 루트로 하는 서브트리의 정점의 수가 저장됩니다.

import sys

sys.setrecursionlimit(10 ** 9)
input = sys.stdin.readline

N, R, Q = map(int, input().split())

tree = [[] for _ in range(N + 1)]

for _ in range(N - 1):
    U, V = map(int, input().split())
    tree[U].append(V)
    tree[V].append(U)

query = [int(input()) for _ in range(Q)]
size = [0] * (N + 1)

def dfs(now):
    size[now] = 1
    for next in tree[now]:
        if size[next] == 0:
            dfs(next)
            size[now] += size[next]

dfs(R)

for i in query:
    print(size[i])

Categories:

Updated: