15681번: 트리와 쿼리

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

 

임의의 루트가 있는 트리가 에서, 입력 U를 루트로 하는 서브트리의 노드 개수를 구하는 문제이다.

 

아이디어

DFS recursion으로 leaf에 도달했을 때, leaf node의 서브 트리의 노드개수는 1

그리고 leaf의 부모 노드로 돌아가면 자식 노드들 개수를 본인 노드 개수와 합친다.

 

# https://www.acmicpc.net/problem/15681
# Tree, DP

import sys
from collections import defaultdict, deque

sys.setrecursionlimit(10**6)

N, R, Q = map(int, sys.stdin.readline().split())
tree = defaultdict(list)
for _ in range(N-1) :
    u, v = map(int, sys.stdin.readline().split())
    tree[u].append(v)
    tree[v].append(u)

subtree = [0] * (N+1)
visited = [0] * (N+1)

def DFS(node) :
    if visited[node] != 0 :
        return

    visited[node] = 1

    for l in tree[node] :
        if visited[l] :
            continue

        DFS(l)
        subtree[l] += 1

        subtree[node] += subtree[l]

subtree[R] = 1
DFS(R)

for _ in range(Q) :
    u = list(map(int, sys.stdin.readline().split()))[0]
    print(subtree[u])

 

subtree[l] += 1로 본인 노드 개수를 세고

subtree[node] += subtree[l]로 자식 노드 총 수를 더한다.

 

 

본래는 parents list를 따로 구현해서 subtree에 저장하는 방식이었는데,

정점 개수가 많아질 수록 O(N)만큼 시간이 더 걸려서 시간초과가 떴다.

그래서 recursion을 적극적으로 이용하는 방식으로 아이디어를 다시 짰다.