[백준/dfs] 15681번: 트리-쿼리

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

15681번: 트리와 쿼리

 

주어진 그래프를 가지고 자식 노드들을 정리해야지 해당 쿼리에 해당하는 자식노드들을 셀 수 있다고 생각했다. 그래서 먼저 

그래프를 만들고, 자식 노드 갯수를 세는 방식으로 했는데 시간초과가 발생했다. 

 

어차피 루트 노드가 주어지므로, 루트 노드를 중심으로 뻗어나가면 자연스럽게 한 노드에 대한 자식노드가 정리 되기 때문에 굳이 트리를 만들지 않아도 된다 

 

이후 트리에서 마지막 입력받는 정점을 루트로하는 서브트리의 정점의 수를 출력하면 된다

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

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

참고한 블로그

 

-정답풀이:

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

N,R,Q=map(int,input().split())
graph=[[] for _ in range(N+1)]
size=[0]*(N+1)

for _ in range(N-1):
    a,b=map(int,input().split())
    graph[a].append(b)
    graph[b].append(a)
    
def countPoint(x):
    size[x]=1
    for i in graph[x]:
        if not size[i]:
            countPoint(i)
            size[x]+=size[i]
countPoint(R)

for _ in range(Q):
    y=int(input())
    print(size[y])

 

-주의할점

graph=[[]*(N+1)]로 그래프를 초기화 했더니 IndexError가 발생했다. 찾아보니

이렇게 하면 graph=[[]]<- 이렇게 생겨서 값을 추가 할 수있는 인덱스가 없다. 

graph=[[] for _ in range(N+1)] 방식으로 작성하자!