[백준 15681] 트리와 쿼리

문제

간선에 가중치와 방향성이 없는 임의의 루트 있는 트리가 주어졌을 때, 아래의 쿼리에 답해보도록 하자.

• 정점 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줄에 걸쳐 각 쿼리의 답을 정수 하나로 출력한다.

---

풀이

 

문제를 보고 바로 떠올랐던 풀이는 BFS였다

입력받은 노드간의 인접 리스트를 형성하고 자식 노드의 정보를 저장하는 부모 노드 리스트를 만들었다

# 노드 인접리스트 생성
for _ in range(N-1):
    n1,n2 = map(int,input().split())
    Rel[n1].append(n2)
    Rel[n2].append(n1)

# BFS로 트리의 부모자식 관계 형성
q.append(R)
visit[R] = True
while len(q) != 0:
    current = q.popleft()
    for node in Rel[current]:
        if visit[node] == False:
            Tree[current].append(node)
            visit[node] = True
            q.append(node)

 

그리고 형성된 부모 노드 리스트를 사용하여 입력받은 쿼리가 루트가 되는 서브트리의 정점의 수를 BFS로 카운트하고 출력했다

 

코드를 작성하면서도 시간초과가 날 것을 예상했지만 50퍼까지 붙길래 혹시나 했다

 

시간을 줄일 방법을 생각했고 문제 하단의 알고리즘 분류에서 "트리에서의 다이나믹 프로그래밍" 이라는 힌트를 얻고 DFS와 DP를 사용하기로 했다

 

결국 bottom-up으로 최하단 서브트리부터 탐색하며

서브트리 정점 개수 = 루트노드의 자식노드들의 정점 합

공식을 적용하였다

 

# DFS와 DP로 서브트리의 정점 수 저장
def DFS(treeNode):
    for node in Tree[treeNode]:
        if subTree[node] == 0:
            DFS(node)
        subTree[treeNode] += subTree[node]
    subTree[treeNode] += 1

for idx in range(1,N+1):
    if subTree[idx] == 0:
        DFS(idx)

 

인덱스 = 루트노드 로 설정하여 subtree에 각 서브트리의 정점 개수를 저장

subTree[node] == 0 일때만 DFS를 돌려 한번 방문하여 정점의 개수를 저장했던 노드는 접근을 제한

서브트리 정점 개수 = 루트노드의 자식노드들의 정점 합 을 적용하여 저장되어있는 자식 노드들의 정점 수 + 루트정점(+1) 으로 서브트리 정점 수 계산