알고리즘/BOJ

[BOJ/python]15681번 트리와 쿼리

frog-in-well 2022. 7. 6. 22:57

백준 15681번 트리와 쿼리 파이썬

문제 링크 - https://www.acmicpc.net/problem/15681

 

15681번: 트리와 쿼리

트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105) 이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)

www.acmicpc.net

반례를 찾는 것이 중요하다. 문제의 로직을 세울 수 있는 것만큼이나 중요한 것 같다. 쉬운 문제였지만 쉬운 반례를 찾지 못해 한참을 헤맸다.

1. 코드

#트리와 쿼리 (tree + dp ) 220503
#dp[node] = 1 이후에 
#if len(tree[node]==1) : return 했는데 틀렸다. 반례 무엇이지? 
#반례 = > 노드 2개인 경우! if len(tree[node]==1) and node!=R : return 이렇게 하면 풀린다!!

import sys
sys.setrecursionlimit(10**4)
N, R, Q = map(int,input().split())

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

dp = [-1 for _ in range(N+1)]

def find(node):
     
    dp[node] = 1
    
    for next in tree[node] :
        if dp[next] == -1 :
            dp[node] += find(next)
    
    return dp[node]

find(R) 
sol = []
for _ in range(Q) :
    U = int(input())
    sol.append(dp[U])
    
for i in range(Q):
    print(sol[i])