[백준] 15681 트리와 쿼리 (파이썬 python)

 

 

문제 링크

 

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

15681번: 트리와 쿼리

 

 

문제 풀이

 

root에 따라 서브쿼리 노드의 개수가 달라지기 때문에

루트노드부터 순차적으로 서브쿼리를 구해나가야 합니다.

특정 노드가 루트가 될 때 그 하위에 포함된 노드들에 대한 서브 쿼리를 구해나가야 하는데

재귀를 통해 각 노드들이 루트 노드가 될 때 하위 노드의 개수를 가장 말단 노드까지 구해줍니다.

 

예를 들어 위 상황 같은 경우는 5 하위의 노드 수를 구하기 위해

4 하위의 노드 수와 6 하위의 노드 수를 구해야 합니다.

4 하위의 노드 수를 구하기 위해서는 3 하위의 노드 수를 구해야 합니다.

이렇게 루트 노드가 될 수 없는 말단 노드전까지 재귀를 타고 들어가 하위 노드 수를 구한 후

밖으로 빠져나오며 하위 노드 수를 더해주는 것입니다.

3 하위의 노드 수가 2이니 본인을 포함해서 정점의 수는 3이 될 것입니다.

4 하위의 노드 수는 3이니 본인을 포함하여 정점의 수가 4가 나올 것이며

5 하위의 노드 수는 4 하위의 노드 수 + 6 하위의 노드 수 + 1 (본인)이 될 것입니다.

 

import sys
input = sys.stdin.readline
sys.setrecursionlimit(10**9)

def get_count(root):
    count[root] = 1
    for node in graph[root]:
        if not count[node]:
            get_count(node)
            count[root] += count[node]


# n : 정점 수, r : 루트 번호, q : 쿼리 수
n, r, q = map(int, input().split())
graph = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    a, b = map(int, input().split())
    graph[a].append(b)
    graph[b].append(a)

count = [0] * (n + 1)
get_count(r)

for _ in range(q):
    u = int(input())
    print(count[u])