[Algorithm] cpp 백준 1763번: 트리 색칠

문제

링크: https://www.acmicpc.net/problem/1763
요약: 루트가 있는 트리에서 노드 i를 F[i]번째로 색칠하면 비용이 C[i]·F[i]. 부모가 먼저 색칠되어야 하며(루트 선행), 전체 비용 합을 최소화하는 색칠 순서를 구한다.
제한/스펙: N ≤ 1000, 1 ≤ R ≤ N. 간선은 부모 U, 자식 V가 직접 주어진다.

입력/출력

예제 1

접근 개요

핵심 관찰: 어떤 노드의 서브트리를 하나의 블록으로 봤을 때, 블록의 현재 누적가중치(sum)와 누적크기(size)를 유지하면, 두 블록 A,B를 부모 쪽으로 올려 쌓는 순서를 바꿔도 비용 차이는 (sumA/sizeA)와 (sumB/sizeB)의 비교로 결정된다. 비율이 큰 블록을 먼저 올릴수록 전체 비용이 감소한다.
그리디 전략: 매 단계 방문되지 않은 정점들 중 비율 sum/size가 최대인 정점 x를 고르고, 아직 방문 전인 가장 가까운 조상 dad로 병합한다. 비용은 sum[x]·size[dad]만큼 증가한다. 병합 후 dad의 (sum,size)에 x의 값을 누적한다.
구현 포인트: 비율 비교는 부동소수점 오차를 피하기 위해 교차곱(sumA·sizeB vs sumB·sizeA)을 정수로 비교하고, 조상 찾기는 “아직 색칠 전인 가장 가까운 조상”으로 경로 압축해 가속한다.

알고리즘

입력으로 부모-자식이 직접 주어지므로 parent[v]=u 배열을 구성하고, 루트 R에 대해서는 parent[R]=R로 둔다.
각 정점은 독립 블록으로 시작: sum[i]=C[i], size[i]=1, visited[i]=false.
N-1회 반복(루트를 제외한 모든 정점을 한 번씩 병합):
- 비방문 정점 i(루트 제외) 중에서 sum[i]/size[i] 비율이 최대인 best를 선출.
- dad = findNearestUnvisitedAncestor(parent[best])를 구해 비용 ans += sum[best] * size[dad]를 더한 뒤, dad 블록에 best를 병합: sum[dad]+=sum[best], size[dad]+=size[best], visited[best]=true.
- findNearestUnvisitedAncestor는 “이미 병합 끝난 정점은 부모로 건너뛰는” 경로 압축으로 구현.
마지막으로 루트 블록 자체를 색칠하는 비용 + sum[R]를 더한다.

올바름(스케치)

두 블록 A,B를 쌓는 순서 비교에서 비용 차이는 (sumA/sizeA) ≥ (sumB/sizeB)일 때 A를 먼저 배치하는 편이 손해가 없다(교환 논증). 이 국소 최적 선택을 전체 과정에 반복 적용하면 전역 최적 순서를 얻는다.
부모 선행 제약은 “아직 방문 전인 가장 가까운 조상으로만 병합” 규칙으로 보존되며, 한 번 병합된 블록은 더 큰 블록으로만 커져 비율 비교의 일관성이 유지된다.

복잡도

매 단계 O(N) 선출 × (N-1) 단계 = O(N^2). N ≤ 1000이므로 충분히 통과한다. 추가 메모리 O(N).

구현 (C++)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
// 더 많은 정보는 42jerrykim.github.io 에서 확인하세요.
#include <bits/stdc++.h>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int N, R;
    if (!(cin >> N >> R)) return 0;

    vector<long long> weight(N + 1);
    for (int i = 1; i <= N; ++i) cin >> weight[i];

    vector<int> parent(N + 1, 0);
    for (int i = 0; i < N - 1; ++i) {
        int U, V; // U is parent of V
        cin >> U >> V;
        parent[V] = U;
    }
    parent[R] = R;

    vector<long long> groupWeight = weight; // 누적 가중치
    vector<int> groupSize(N + 1, 1);        // 누적 크기
    vector<char> visited(N + 1, 0);         // 병합 완료 표시

    function<int(int)> findNearestUnvisitedAncestor = [&](int x) -> int {
        if (!visited[x]) return x;
        return parent[x] = findNearestUnvisitedAncestor(parent[x]);
    };

    long long answer = 0;

    // N-1번: 루트를 제외한 모든 정점을 한 번씩 선택해 위로 병합
    for (int step = 0; step < N - 1; ++step) {
        int best = -1;
        for (int i = 1; i <= N; ++i) {
            if (i == R || visited[i]) continue;
            if (best == -1) {
                best = i;
            } else {
                // 최대 ratio: groupWeight[i] / groupSize[i]
                __int128 lhs = (__int128)groupWeight[i] * groupSize[best];
                __int128 rhs = (__int128)groupWeight[best] * groupSize[i];
                if (lhs > rhs) best = i;
            }
        }

        int dad = findNearestUnvisitedAncestor(parent[best]);
        visited[best] = 1;

        // 이번에 best 블록을 dad 위로 올려 쌓으면서 발생하는 비용
        answer += groupWeight[best] * (long long)groupSize[dad];

        // dad 블록에 병합
        groupWeight[dad] += groupWeight[best];
        groupSize[dad] += groupSize[best];
        parent[best] = dad; // 경로 압축 일관성
    }

    // 마지막에 루트가 칠해질 때 드는 비용
    answer += groupWeight[R];

    cout << answer << '\n';
    return 0;
}