[백준][Python] 15681번 트리와 쿼리

15681번: 트리와 쿼리

 

 

코드

import sys
sys.setrecursionlimit(10**9)
def tree(node,par):
	for i in graph[node]:
		if i != par:
			child[node].append(i)
			parent[i] = node
			tree(i,node)
def count(node):
	size[node] = 1
	for i in child[node]:
		count(i)
		size[node] += size[i]
n,r,q = map(int,input().split())
graph = [[] for _ in range(n+1)]
child = [[] for _ in range(n+1)]
parent = [-1] * (n+1)
size = [0] * (n+1)
for _ in range(n-1):
	u,v = map(int,sys.stdin.readline().split())
	graph[u].append(v)
	graph[v].append(u)
tree(r,-1)
count(r)

for _ in range(q):
	u = int(sys.stdin.readline())
	print(size[u])

 

설명

문제 아래 나와있는 힌트를 보고 풀었다.

먼저 주어진 입력에 따라 그래프를 만들고 tree 함수에서 이 그래프를 이용해서 각 노드의 자식과 부모 노드를 각각 저장한다.

count함수는 dfs를 사용하여 정점의 개수를 구한다.