[백준/Python3] 15681 트리와 쿼리

춘배컵에 참여해서 뱃지를 받았어요 ㅎㅎ

✈Problem

트리와 쿼리
간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.
만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.

Solution

문제 하단부의 힌트를 참고해서 풀이했다. 먼저 루트를 기준으로 트리 구조를 만든다고 생각하고 각 간선의 자식 노드들을 구해 child에 저장해줬다. 트리 구조로 만들었다면 각 노드의 서브트리의 개수를 구해서 size[node]에 저장해주었다. 한번 각 노드의 서브트의 개수는 각 노드가 루트일 때의 서브트리의 개수와 동일하므로 각 쿼리의 해답은 size[node]가 된다.

Python Code

from sys import stdin, setrecursionlimit

setrecursionlimit(1000000000)

def make_tree(current_node, parent):
    for node in connect[current_node]:
        if node != parent:
            child[current_node].append(node)
            make_tree(node, current_node)

def count_subtree_nodes(current_node):
    size[current_node] = 1
    for node in child[current_node]:
        count_subtree_nodes(node)
        size[current_node] += size[node]

N, R, Q = map(int, stdin.readline().split())
connect = [[] for _ in range(N + 1)]

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

child = [[] for _ in range(N + 1)]
make_tree(R, -1)
size = [0] * (N + 1)
count_subtree_nodes(R)
for _ in range(Q):
    root = int(stdin.readline())
    print(size[root])