BOJ 4002, APIO 2012 닌자배치(Dispatching)

• 풀이
  • 문제 해석
  • Algorithm
  • Small-to-Large Technique
• 코드

풀이

문제 해석

번역이 상당히 난해하게 작성되어 있다. 요점을 정리하자면

• 루트 있는 트리가 주어진다.
• 각 노드에는 $L_i, C_i$ 값이 주어진다.
• 트리의 어떤 노드 하나 (매니저) 를 고른다.
• 그 트리를 루트로 하는 서브트리에서 몇 개의 노드(닌자)를 고른다. 매니저를 닌자로 삼을 수 있다.
• 이때, 고른 닌자 (매니저를 포함하지 않는다) 의 $C_i$ 값의 합이 $M$을 넘어서는 안 된다.
• 나의 점수 는, (매니저의 $L$ 값) * (고른 닌자의 수) 로 정의된다.
• 점수를 최대화하라.

Algorithm

이 문제를 해결하기 위해 약간 Naive한 솔루션을 생각해 보자. 각 노드마다, 내가 매니저일 때 고르는 최선의 닌자 조합 을 가지고 있는다고 가정하자. 내가 매니저일 때는 나의 서브트리에서 닌자들을 골라야 하고, 이때 각 닌자는 모두 똑같이 취급되므로 무조건 월급 $C_i$ 가 낮은 값싼 닌자를 많이 배치할수록 이득이다. 이 배치를 리프 노드부터 갱신한다고 생각하자. 이때, 닌자 $x$에 대한 최선닌자조합을 $S_x$ 라는 집합으로 정의하자.

관찰 $x$의 ancestor node $p$에 대해, $S_x$에 포함되지 않는 노드는 $S_p$에 포함되지 않는다.

증명 $S_p$를 어떻게 만들지 생각해 보자. 먼저, $S_p$를 만드는 가장 쉬운 방법 중 하나는, 자기 자식 노드 $v_1, \dots v_k$ 에 대해 $S_{v_i}$ 들을 모두 합집합한 다음, 비용이 $M$보다 작거나 같아질 때까지 비싼 닌자를 계속 자르는 방법이 있다. 이 방법이 올바른 $S_p$를 보장함은 어렵지 않게 보여진다.

이 관찰로부터 얻어지는 알고리즘의 시간 복잡도를 생각해 보자. 각 노드마다, 자기 자식 노드의 모든 set을 합쳐야 하고, 필요하다면 닌자들을 해고해야 한다. 이때 $S_x$를 set이나 priority queue와 같은 자료구조를 이용하여 관리한다면, 최대 $n$번 합집합 연산을 해야 하며, 각 집합의 크기가 최대 $n$까지 갈 수 있고, 원소 하나를 움직이는 데 $O(\log n)$ 시간이 들 것이므로 $O(n^2 \log n)$ 시간이 소요된다.

Merge sort tree를 만드는 때와 비슷하게, set을 쓰지 않고 vector 같은것을 정렬해서 가지고 있으면서 merge하면 $O(n^2)$ 시간에도 어떻게 가능할 것 같기도 하다. 이 방법은 구현해 보지 않아서 복잡도에는 자신이 없다.

Small-to-Large Technique

아주 작은 부분을 바꿈으로써 시간 복잡도를 줄이는 테크닉인데, 굉장히 아름답다. 요지는, 위 알고리즘을 그대로 수행하되, set을 합칠 때 항상 작은 집합에 있는 원소들을 큰 집합으로 옮기는 방법으로만 합치는 것이다. 여러 개를 합칠 떄는 가장 큰 것을 기준으로 나머지를 거기에 집어넣기로 하자.

이 방법이 복잡도를 줄여 줌을 증명하자.

증명 각 닌자 $x$마다, $x$가 포함된 집합이 있을 것이다. (집합에서 한 번 빠진 닌자는 사라지는 것이므로, 그걸 새로운 집합처럼 생각하자) 위 알고리즘의 수행 과정을 손으로 그려 보면, 각 원소마다 모든 집합들을 통틀어 한 번만 있으면 된다 (다 계산한 다음 자식 노드의 $S_v$ 들을 모두 지워버려도 되므로).
따라서, $x$가 포함된 집합 $T(x)$ 를 항상 identify 할 수 있다.
이때, $T(x)$의 크기를 반환하는 함수 $f(x)$를 생각하자. $f(x)$는, 즉, 지금 이 순간 x가 포함되어 있는 집합의 크기 라고 정의한다. 원소 $x$가 어떤 집합 $A$에서 다른 집합 $B$로 이동한다면, 즉 $T(x)$가 $A$에서 $B$로 변경된다면 어떤 일이 일어나는지 생각해 보자. 만약 $A$가 $B$보다 작았다면, $x$ 하나만 옮기는게 아니라 $A$를 통째로 들어다가 $B$로 옮길 때맏 $x$가 움직이므로 $T(x)$는 $\abs{A}$에서 $\abs{A} + \abs{B}$로 바뀌게 되어, 적어도 2배 이상 늘어난다. 반면 만약 $A$가 $B$보다 컸다면, $A$와 $B$를 합치더라도 $x$는 움직이지 않고 $B$에 있었던 원소들이 이쪽으로 움직일 것이다. 즉, $x$를 relocate 하는 연산이 일어난다면 반드시 $f(x)$가 2배 이상 증가한다는 것을 알 수 있다.

이때, 당연하지만 중요한 사실은 당연히 $f(x)$는 $n$을 넘을 수 없다는 점이다. 따라서, 이 사실들을 종합하면, $x$는 최대 $O(\log n)$ 번만 이동할 수 있다는 결과를 얻을 수 있다.

따라서, 이 테크닉을 적용하면, 모든 원소는 최대 $O(\log n)$ 번 움직이고, 움직일때마다 $O(\log n)$ 시간이 소모되므로 전체 알고리즘의 시간 복잡도가 $O(n \log^2 n)$으로 크게 줄어든다. 이 테크닉을 smaller-to-larger 같은 이름으로 부르기도 하는데, 요점은 약간 DSU에서 union by size 할 때 복잡도 증명할 때와 같다.

구현도 굉장히 쉽게 할 수 있다.

코드

#include <bits/stdc++.h>
#pragma GCC optimize("O3")
#pragma GCC optimize("Ofast")
#pragma GCC target("avx,avx2,fma")
#define ll long long
#define int ll
#define eps 1e-7
#define all(x) ((x).begin()),((x).end())
#define usecppio ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);
using namespace std;
using pii = pair<int, int>;

struct PQueue
{
    priority_queue <int> pq;
    int total = 0;
    void push(int x)
    {
        pq.push(x);
        total += x;
    }
    int peek() {return pq.top();}
    int pop()
    {
        int x = pq.top();
        total -= x;
        pq.pop();
        return x;
    }
    int size() {return (int)(pq.size());}
    bool empty() {return pq.empty();}
};

void merge(PQueue &dest, PQueue &src)
{
    if (dest.size() < src.size())
        swap(dest, src);
    while (!src.empty())
        dest.push(src.pop());
}

int n, m, ans;
vector <int> T[101010];
vector <PQueue> PQS(101010);
int lead[101010], cost[101010];

void dfs(int r, int p)
{
    PQS[r].push(cost[r]);
    for (int c : T[r])
        if (c != p)
            dfs(c, r);
    for (int c : T[r])
        if (c != p)
            merge(PQS[r], PQS[c]);
    while (PQS[r].total > m)
        PQS[r].pop();
    ans = max(ans, lead[r] * PQS[r].size());
}

int32_t main()
{
    usecppio
    cin >> n >> m;
    for (int i = 1; i <= n; i++)
    {
        int b, c, l; cin >> b >> c >> l;
        T[b].push_back(i);
        T[i].push_back(b);
        lead[i] = l;
        cost[i] = c;
    }
    cost[0] = INT_MAX;
    dfs(0, -1);
    cout << ans << '\n';
}