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

문제

www.acmicpc.net/problem/15681

15681번: 트리와 쿼리

 

풀이

문제가 어려워 보이지만, 간단히 말하면 트리에서 마지막 입력받는 정점을 루트로하는 서브트리의 정점의 수를 출력하면 된다. 

dfs로 탐색하면서 각 정점에 대한 서브트리의 수를 count 배열에 담아준다.

count 배열로 방문여부 표시도 같이 해줄 수 있다.

소스코드

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

def countPoint(x):
    count[x]=1
    for i in tree[x]:
        if not count[i]:
            countPoint(i)
            count[x] += count[i]



n, r, q = map(int, input().split())
tree = [[] for _ in range(n+1)]
count = [0]*(n+1)

for i in range(n-1):
    a, b = map(int, input().split())
    tree[a].append(b)
    tree[b].append(a)

countPoint(r)

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