[백준 32072번] 트리 뽑아내기

문제 링크

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

 

문제 요약

문제에서 정의한 뽑아내기 연산을 수행할 때마다 루트에 있는 가중치의 값을 출력하는 문제이다. 뽑아내기 연산을 수행할 때는 특별한 경로를 찾아야 하는데, 한 정점에서 이동할 자식을 결정할 때는 직계 자식들만 비교하면 된다는 것에 유의하자. (한 정점의 자손들 중 가장 작은 가중치를 가지는 정점으로 이동하는 경로를 찾아야 하는 것으로 문제를 잘못 이해해 시간을 너무 많이 날렸다...🤦‍♂️)

 

문제 풀이

뽑아내기 연산을 수행한다고 하지 말고, 루트 정점에 있어야 하는 정점들을 순회한다고 접근하면 문제를 조금 더 쉽게 풀 수 있다. 그렇다면 어떤 순서대로 정점을 방문해야 할까?

 

뽑아내기 연산을 할 때, 특별한 경로에 있는 정점들 중 하나인 정점 \(V\)와 자식 정점 \(K\)가 있다고 가정해 보자. 그리고 \(V\)와 같은 레벨에 있는 정점 \(U\)도 있다. 뽑아내기 연산 후 \(V\)의 가중치는 \(K\)의 가중치가 되기 때문에, 다음 뽑아내기 연산에 \(U\)가 포함되려면 \(U\)의 가중치는 \(K\)의 가중치보다 작아야 한다. 만약 \(U\)의 가중치가 \(K\)의 가중치보다 크다면, 다음 뽑아내기 연산에 포함될 정점은 \(K\)가 된다.

 

이 사실을 확장하면, 트리의 순서를 유지하면서 가중치가 가장 작은 정점들부터 방문하는 것이 뽑아내기 연산을 \(N\)번 수행했을 때의 결과와 같다는 것을 알 수 있다.

 

소스 코드

#include <stdio.h>
#include <vector>
#include <queue>
#include <functional>

using namespace std;

int n, p, A[300010], H[300010];
vector<int> node[300010];
priority_queue<int, vector<int>, greater<int>> pq;

int main() {
    scanf("%d", &n);
    
    for (int i = 2 ; i <= n ; ++i) {
        scanf("%d", &p);
        node[p].push_back(i);
    }
    
    for (int i = 1 ; i <= n ; ++i) {
        scanf("%d", &A[i]);
        H[A[i]] = i;
    }
    
    pq.push(A[1]);
    
    while (!pq.empty()) {
        int val = pq.top();
        int here = H[val];
        
        pq.pop();
        
        printf("%d\n", val);
        
        for (int there: node[here]) {
            pq.push(A[there]);
        }
    }
}