백준 15681번 : 트리와 쿼리 in Python

15681번: 트리와 쿼리

 

www.acmicpc.net

이 문제는 루트 있는 트리가 주어졌을 때 임의의 정점을 루트로 하는 서브트리에 속한 정점의 수를 출력하는 문제이다.

우선 트리 내의 서열을 고려한 그래프 탐색을 위해서는 DFS를 이용해야 한다. 왜냐하면 DFS의 경우 함수를 재귀적으로 호출할 때 자식 노드와 부모 노드가 명확히 결정되기 때문이다. 예를 들어 루트 노드에 대해 DFS를 호출하면 루트 노드의 자식 노드에 대해 DFS를 재귀적으로 호출하고, 자식 노드에 대한 DFS 내에서는 그 자식 노드의 자식 노드들에 대해 DFS를 재귀적으로 호출한다. 

그리고 이 문제에서 얻고자 하는 것은 임의의 노드들에 대해 해당 노드가 루트가 되는 서브트리의 정점의 개수인데, 주어지는 노드의 개수가 10^5개까지 가능하기 때문에 노드가 주어질 때마다 서브트리를 조사하기에는 시간이 오래 걸린다. 그래서 노드가 주어질 때마다 서브트리의 정점 개수를 바로 출력하는 것이 효율적이므로, 이를 미리 저장해 두는 것이 효율적일 것이다. 즉, dynamic programming을 이용한다. 

이를 위해 각 정점마다 해당 정점이 루트일 때 서브트리의 정점의 개수를 저장하는 배열인 subtree_num을 만들고, 각 원소를 1로 초기화한다. 그 다음 DFS 함수 내에서 현재 정점과 인접한 정점들에 대해 DFS 함수를 재귀적으로 호출한 뒤 각 정점들이 루트가 되는 서브트리의 정점의 개수를 현재 정점이 루트가 되는 서브트리의 정점의 개수에 더해준다. 그러면 인접한 정점들에 대해 DFS 함수를 호출한 뒤에는 그 정점들에 대한 서브트리를 모두 조사한 뒤이므로 그 정점들에 대해 서브트리의 정점의 개수를 정상적으로 구한 상태이므로, 이러한 알고리즘은 정상적으로 현재 노드의 서브트리의 정점의 개수를 구한다.

이를 코드로 구현하면 다음과 같다.

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

# DFS 함수
def DFS(currNode):
    if visited[currNode]: return
    visited[currNode] = True
    for v in graph[currNode]:
        if not visited[v]:
            DFS(v)
            # 서브트리의 정점의 개수에 자식 노드의 서브트리의 정점의 개수를 더해준다.
            subtree_num[currNode] += subtree_num[v]
N, R, Q = map(int, input().split())
graph = [[] for _ in range(N+1)]
for _ in range(N-1):
    U, V = map(int, input().split())
    graph[U].append(V)
    graph[V].append(U)
visited = [False for _ in range(N+1)]
subtree_num = [1 for _ in range(N+1)]
# 루트 노드에 대해 DFS를 수행한다.
DFS(R)
for _ in range(Q):
    U = int(input())
    print(subtree_num[U])