백준 15681 : 트리와 쿼리 (Python)

본문 링크

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


def MakeTree(start):

    distance[start]=1

    for i in Tree[start]:
        if distance[i]==0:
            MakeTree(i)
            distance[start]+=distance[i]

N,R,Q=map(int,input().split())

Tree=[ [] for _ in range(N+1) ]
distance=[0]*(N+1)

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

MakeTree(R)
for i in range(Q):
    c=int(input())
    print(distance[c])

📌 어떻게 접근할 것인가?

다행히 문제에서 힌트가있어 설명이 아주 잘 되어있었다.
만약 정점의 수를 매번 쿼리마다 계산하면 시간초과가 난다.
따라서 distance 리스트를 선언하여 한번이 탐색으로 모든 정점의 수를 계산한다.

이런 트리가 있다고 해보자.
먼저 리프노드에 도착하면 리프노드의 값을 1로 초기화 한다. 이는 $DFS$ 탐색을 통해 구현가능하다.
이후 백트래킹과 유사한 방법으로 재귀 호출 이후에 distance[start]+=distance[i] 을 해준다.
즉 자식노드의 값을 부모노드에 모두 더해준다.
이 과정을 통해 문제에서 구하고자 하는 방법을 바로 구할수 있다.

• 정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.

또한 자기자신도 정점이기 때문에 자기자신에는 항상 값을 1 넣어준다.
코드를 보면 조금더 이해가 수월할 것이다.
$Tree$ 를 선언해주고 $N-1$번만큼 정점과 간선을 입력받아 연결해 준다.
이후 루트노드 $R$을 시작점으로 재귀함수를 실행시켜준다.

distance 을 트리에서의 동적계획법으로 만들었으므로
$Q$번의 쿼리가 주어질때 $c$ 를 입력받아서 distance[c] 을 출력하면 $O(1)$ 의 시간복잡도로 문제 해결이 가능하다.

트리 DP의 기초 문제이니 꼭 이해하도록 하자.