백준 15681 - 트리와 쿼리 (파이썬)

난이도

골드 5

유형

트리에서의 다이나믹 프로그래밍

DFS

해결

트리DP는 처음 해보았지만 문제에서 뭔가 느낌이 왔다. DFS를 쓰고 방문처리를 해주면 될 것 같았다.

dfs 사용시 메모리 초과와 재귀횟수 제한 설정에 유의해줘야 하는 문제였다.

코드

from sys import stdin,setrecursionlimit
from collections import defaultdict
input=stdin.readline
setrecursionlimit(1000000)

n,r,q=map(int,input().split())
tree=[[] for i in range(n+1)] #트리
dp=[1]*(n+1) # 정답
vi=[0]*(n+1) # 방문처리

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

# 루트 노드는 방문처리

def dfs(node):
    global dp
    vi[node]=1
    for i in tree[node]:
        if vi[i]==0:
            dfs(i)
            dp[node]+=dp[i]
    
    
dfs(r)

for i in range(q):
    print(dp[int(input())])

좋은 / 비슷한 문제