[골드 5] 백준 15681 - 트리와 쿼리 (파이썬)

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

풀이

• 루트가 R인 트리가 주어짐
• 각 쿼리는 정점 U를 루트로 하는 서브트리의 정점 개수를 묻는다
• 여러 쿼리에 대해 빠르게 답변해야 하므로 DFS를 활용한 서브트리 크기 미리 계산이 핵심

adj = [[] for _ in range(N)]

adj[u]에 연결된 정점들 v를 저장하여 무방향 트리를 구성한다.

D = [0] * N

D[u]는 u를 루트로 하는 서브트리의 정점 수를 나타낸다.

def dfs(u):
    visit[u] = True
    D[u] = 1  # 자기 자신 포함

    for v in adj[u]:
        if not visit[v]:
            dfs(v)
            D[u] += D[v]

DFS를 통한 서브트리 크기를 계산하며, 재귀적으로 자식 노드의 서브트리 크기를 더해준다.

한번의 DFS로 모든 정점의 서브트리 크기 계산이 가능하다.

for _ in range(Q):
    u = int(input())
    u -= 1
    print(D[u])

이미 계산해둔 D[u] 값을 출력하면 되므로 O(1) 시간을 가진다.

코드

# 백준 15681 - 트리와 쿼리
# 분류 : 다이나믹 프로그래밍

import sys
sys.setrecursionlimit(100000000)
input = sys.stdin.readline

N, R, Q = map(int, input().split())
R -= 1
adj = [[] for _ in range(N)]

for _ in range(N - 1) :
    u, v = map(int, input().split())
    u -= 1
    v -= 1

    adj[u].append(v)
    adj[v].append(u)

visit = [False] * N
D = [0] * N

def dfs(u) :
    visit[u] = True
    D[u] = 1

    for v in adj[u] :
        if not visit[v] :
            dfs(v)
            D[u] += D[v]

dfs(R)

for _ in range(Q) :
    u = int(input())
    u -= 1

    print(D[u])