백준[ 15681 | Python ] 트리와 쿼리

문제

https://www.acmicpc.net/problem/15681

풀이

그래프 정보와 루트 정보를 주고, 쿼리에 따른 서브트리의 정점 수를 새는 문제다.
먼저 루트로부터 주어진 그래프를 DFS로 탐색한다.
탐색을 수행하면서 들리는 정점마다 개수를 새어준다. 정점의 개수는 들리는 다음 노드에 저장한다.
마지막으로 쿼리에 따라 노드 개수를 출력한다.

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

n, r, q = map(int,input().split())
graph = [[] for _ in range(n+1)]
v = [0]*(n+1)

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

query = []
for i in range(q):
    query.append(int(input().rstrip()))

def dfs(start):
    v[start] = 1
    for i in graph[start]:
        if not v[i]:
            dfs(i)
            v[start]+=v[i]

dfs(r)
for i in query:
    print(v[i])

생각 & 정리

DFS와 동적프로그래밍을 통해 해결할 수 있는 문제였다.
처음 문제를 봤을 때 쿼리의 수 만큼 DFS탐색을 수행하려 했는데, 알고리즘 분류를 훑고나서 깨닫게 되었다.
정확하게 이해를 하지 못한탓인지, BFS문제들 보다는 DFS 문제들이 좀 더 다양하게 느껴진다.