백준 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])