Category
알고리즘  DFS/BFS

[15681] 트리와 쿼리

by Hwang

문제

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

출력

Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

풀이

정점 U를 루트로 하는 서브트리의 정점의 수를 매번 구하는 방식은 효율적이라 볼 수 없습니다.

가장 최상위 정점인 R의 서브트리의 속한 정점의 수를 구해주었고, 이 과정에서 거친 모든 정점에서의 서브트리의 정점의 수를 배열에 저장해주었습니다.

이렇게 미리 구한 값을 이용해서 쿼리에 답을 해주는 방식으로 시간 초과를 피해갈 수 있습니다.

import sys

sys.setrecursionlimit(10 ** 5)
N, R, Q = map(int, sys.stdin.readline().split())
tree = [[] for _ in range(N + 1)]
visited = [False] * (N + 1)
dp = [0] * (N + 1)

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

def dfs(current_node):
    visited[current_node] = True
    count = 1
    for node in tree[current_node]:
        if not visited[node]:
            count += dfs(node)
    dp[current_node] = count
    return count

dfs(R)
for _ in range(Q):
    U = int(sys.stdin.readline())
    print(dp[U])


 Comment