트리와 쿼리
문제
간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.
정점 U를 루트로 하는 서브트리에 속한 정점의 수를 출력한다.
만약 이 문제를 해결하는 데에 어려움이 있다면, 하단의 힌트에 첨부한 문서를 참고하자.
입력
트리의 정점의 수 N과 루트의 번호 R, 쿼리의 수 Q가 주어진다. (2 ≤ N ≤ 105, 1 ≤ R ≤ N, 1 ≤ Q ≤ 105)
이어 N-1줄에 걸쳐, U V의 형태로 트리에 속한 간선의 정보가 주어진다. (1 ≤ U, V ≤ N, U ≠ V)
이는 U와 V를 양 끝점으로 하는 간선이 트리에 속함을 의미한다.
이어 Q줄에 걸쳐, 문제에 설명한 U가 하나씩 주어진다. (1 ≤ U ≤ N)
입력으로 주어지는 트리는 항상 올바른 트리임이 보장된다.
출력
Q줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.
힌트를 읽다가 머리가 더 아파졌었다.글이 너무 많더라고
그래프 탐색과 DP를 섞어서 풀면 좋을 것 같다는 힌트를 얻고 문제에 접근 하였다. 이전에 많은 그래프 탐색 문제를 풀면서 그래프 간의 거리를 구할 때 종종 쓰곤 했던 방식과 유사한 것 같았고, 자신의 부모 노드를 구분하는 것, 그리고 루트 노드를 기억하는 것이 중요해보이는 문제였다.
구현(DFS + DP)
루트 노드부터 DFS를 수행한다.
방문한 노드를 방문처리 하고, 거쳐온 노드의 수를 적어둘 무언가를 만든다.
모든 노드를 방문하며, 자식 노드의 수를 무언가(subtree)에 저장한다.
빨간색이 DP가 사용될 부분이고, 기본적으로 DFS를 기반으로 한다.
코드를 살펴보자
import sys
sys.setrecursionlimit(10**9)
input = sys.stdin.readline
def dfs(root):
subtree[root] = 1
visited[root] = True
for i in tree[root]:
if not visited[i]:
dfs(i)
subtree[root] += subtree[i]
N,R,Q = map(int,input().split())
tree = [[]*(N+1) for _ in range(N+1)]
subtree = [0]*(N+1)
visited = [False] * (N+1)
for i in range(N-1):
a,b = map(int,input().split())
tree[a].append(b)
tree[b].append(a)
dfs(R)
for _ in range(Q):
a = int(input())
print(subtree[a])
시험기간 동안, 알고리즘을 쉬었더니 어떤 식으로 문제를 풀었는지 설명하기도 힘들고, 문제 푸는 것도 굉장히 어려웠다.
