bj15829 트리와 쿼리

문제

• 백준 URL: 트리와 쿼리
• 난이도 골드 5

풀이

import sys
from collections import deque

input = sys.stdin.readline

n,r,Q = map(int,input().split())

adj = [[] for _ in range(n+1)]
for _ in range(n-1):
    u,v = map(int,input().split())
    adj[u].append(v)
    adj[v].append(u)

tree = [[] for _ in range(n+1)]
visited = [0]*(n+1)
order = []

q = deque([r])
visited[r] = 1

while q:
    node = q.popleft()
    order.append(node)
    for nxt in adj[node]:
        if not visited[nxt]:
            visited[nxt] = 1
            tree[node].append(nxt)
            q.append(nxt)

subtree = [0]*(n+1)

for u in reversed(order):
    size = 1
    for v in tree[u]:
        size += subtree[v]
    subtree[u] = size

for _ in range(Q):
    u = int(input())
    print(subtree[u])

간단한 트리 문제이다. 가장 끝에 있는 리프 노드에서부터 서브트리 크기를 계산하는 것이 핵심이다.

n,r,Q = map(int,input().split())

adj = [[] for _ in range(n+1)]
for _ in range(n-1):
    u,v = map(int,input().split())
    adj[u].append(v)
    adj[v].append(u)

tree = [[] for _ in range(n+1)]
visited = [0]*(n+1)
order = []

q = deque([r])
visited[r] = 1

while q:
    node = q.popleft()
    order.append(node)
    for nxt in adj[node]:
        if not visited[nxt]:
            visited[nxt] = 1
            tree[node].append(nxt)
            q.append(nxt)

여기까지는 입력을 받고 트리를 구성해주는 단계이다.
가장 먼저 '인접한 노드를 가진 그래프' 입력받은 다음에 BFS로 '방향성을 가진 트리'로 만들어준다.

subtree = [0]*(n+1)

for u in reversed(order):
    size = 1
    for v in tree[u]:
        size += subtree[v]
    subtree[u] = size

Bottom-Up으로 서브트리의 사이즈를 구한다.
루트 노드로부터 가장 먼 곳에서부터 시작해야 하므로, BFS로 트리를 생성할 때 이 순서를 저장하는 'order' 리스트를 생성해서 이를 reversed() 해주면 가장 끝에 있는 리프 노드에서부터 시작할 수 있다.