BOJ 34113 Migration Plan (JOISC 2024-2025 Day 4 P2)

문제

문제를 트리와 쿼리 형태로 생각하면 다음과 같다.

• 초기에, 트리의 각 점에 $K_i$가 쓰여 있다.
• Relocation: 모든 depth $X_i$ 정점 $v$에 대해, 해당 정점의 $Y_i$ 깊이 조상이 $w$라면 $K_w$에 $K_v$를 더하고, $K_v$를 $0$으로 만든다.
• Immigration: $K_{A_j}$에 $L_j$를 더한다.
• Survey: $K_{B_j}$를 출력한다.

풀이

Subtask 1 (4점)

트리가 스타 그래프 형태이므로, 문제가 매우 쉽다. Relocation은 그냥 모든 $K_i$를 루트로 모으는 것과 같고, Immigration과 Survey도 쉽다.

Subtask 2 (12점)

나이브 $O(NQ)$를 짜면 통과한다.

Subtask 3 (25점)

모든 정점의 깊이가 $20$ 이하라는 조건이 있다.

초기에 존재하는 인구나, Immigration으로 들어오는 인구를 최대 $N+Q$개의 chunk로 나눠 생각할 수 있다. 처음에 각 정점에 chunk가 한 개씩 있고, 나중에 들어오는 Immigration마다 하나의 chunk로 생각하는 식이다. 각 chunk는 인구 수가 다를 수 있을 것이다.

이때 각 chunk는 relocation으로 조상 노드로 이동할 수 있는데, 이 횟수가 최대 $D$번임을 관찰할 수 있다. 따라서, 각 chunk를 나이브하게 올려 줘도 최대 $O((N+Q)D)$ 의 시간 복잡도가 든다는 뜻이다. 이걸 하려면 각 깊이별로 존재하는 chunk만 순회해서 올려 줘야 하고, $chunks[d]$를 현재 깊이 $d$에 존재하는 chunk의 목록으로 정의하게 되면 해결 가능하다.

Subtask 4 (40점)

이 서브태스크에는 $2$번 쿼리가 없다.

일단 가장 자명한 방법은 각 chunk가 어떻게 이동하는지 보는 것이다. 이때, 초기에 깊이가 같은 chunk는 항상 깊이가 같게 이동한다는 사실을 눈치챌 수 있다. 이 사실을 이용해 보자.

depth를 union find와 같이 관리해, 처음에 depth $d$였던 chunk가 현재 depth 얼마에 와 있는지를 계산할 수 있다. survey가 최대 $5$개뿐이므로, 각 쿼리를 $O(n)$에 답할 수 있다. 쿼리 지점을 $x$, 그 깊이를 $d'$라고 하면, $x$의 자손을 모두 둘러보며, 그 깊이를 이용해 현재 $x$에 도달했는지 아닌지를 확인할 수 있다. 이렇게 하면 relocation을 쿼리당 $O(\log n)$ (또는 $O(\log^* n)$)에, survey를 개당 $O(n \log n)$ (또는 $O(\log^* n)$에) 해결할 수 있다.

Subtask 5 (55점)

Subtask 4와 비슷하지만, Immigration이 추가되었다.

만약 어떤 immigration chunk가 이미 존재하는 다른 chunk와 같은 깊이에 추가된다면, 그냥 그 chunk에 합쳐 줘도 별 지장이 없을 것이다. 그렇지 않다면, union find에서 해당 depth를 표현하는 새로운 노드를 하나 만들고 처리해준다. 이렇게 하면 앞과 거의 비슷한 방식으로 문제를 해결할 수 있다.

Subtask 6 (82점)

Subtask 4와 비슷하지만, 3번 쿼리의 개수에 제약이 없다.

Union find를 관리할 때, 각 집합 안에 있는 chunk를 관리하고, merge할 때 small to large의 방식을 이용해 merge한다. 여기서 Euler Tour Trick을 이용해 이 chunk들을 dfs ordering 순서로 관리한다고 생각한다. 이렇게 하면 set에 모든 chunk들을 관리할 수 있고, 대략 $O(Q \log^2 N)$ 정도에 서브태스크를 해결할 수 있겠지만, 제한이 큰 만큼 TLE를 받을 수 있다.

다른 접근 방식은, Union Find에서 각 노드가 합쳐지는 과정을 모방해 Union Find Tree를 만드는 것이다. 이 트리에서 ETT를 돌린 뒤 dfs ordering을 이용해 현재 시점에서 깊이 d'의 노드가 깊이 d로 올라왔는지를 알아낼 수 있다. 이렇게 하면 각 chunk마다 두 가지 dfs ordering을 부여받은 상태가 되는데, 이를 xy평면에 나타내면 직사각형 합 쿼리로 환원할 수 있다. 여기서부터는 오프라인으로 세그 스위핑을 하면 답을 어렵지 않게 찾을 수 있다. 시간 복잡도는 $O((N+Q) \log (N+Q))$이다.

Full Task (100점)

Subtask 5와 6의 방법을 합치면 $O((N+Q) \log (N+Q))$의 시간 복잡도로 문제를 해결할 수 있다.

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

void input();
void makeUFTree();
void sweep();

int main(){
    #ifndef LOCAL
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    #endif

    input();
    makeUFTree();
    sweep();
}

struct Tree{
    int n;
    vector<vector<int> > edges;
    vector<int> in, out;
    int inCnt;

    void init(int _n, vector<pair<int, int> > &vec){
        n = _n;
        edges.resize(n+1);
        for(auto [x, y]: vec){
            edges[x].push_back(y);
            edges[y].push_back(x);
        }
    }

    void inner_dfs(int x, int p=-1){
        in[x] = ++inCnt;
        for(int y: edges[x]){
            if(y==p) continue;
            inner_dfs(y, x);
        }
        out[x] = inCnt;
    }

    void dfsOrder(int root = 1){
        in.resize(n+1);
        out.resize(n+1);
        inner_dfs(root);
    }
};

Tree baseTree, ufTree;

struct Query{
    int type, a, b;
    Query(){}
    Query(int type, int a, int b): type(type), a(a), b(b){}
};

int n, q;
int A[2'000'005];
int depth[2'000'005];
Query queries[2'000'005];

void input(){
    cin >> n;

    vector<pair<int, int> > edges;
    for(int i=2; i<=n; i++){
        int p;
        cin >> p;
        edges.push_back({p, i});
        depth[i] = depth[p] + 1;
    }
    baseTree.init(n, edges);
    baseTree.dfsOrder();

    for(int i=1; i<=n; i++) cin >> A[i];

    cin >> q;
    for(int i=1; i<=q; i++){
        cin >> queries[i].type >> queries[i].a;
        if(queries[i].type <= 2) cin >> queries[i].b;
    }
}

struct Chunk{
    // in1: 기본 트리에서의 dfs ordering
    // in2: UF tree에서의 dfs ordering
    // 만약 UF tree에서 dfs ordering을 계산하기 전이라면 UF 트리에서 해당하는 정점 번호를 넣음
    int in1, in2, val;
    Chunk(){}
    Chunk(int in1, int in2, int val): in1(in1), in2(in2), val(val){}
};

struct SweepEvent{
    int type; // 0: 더하기, 1: 쿼리
    int l, r, idx, w;
    SweepEvent(){}
    SweepEvent(int type, int l, int r, int idx, int w): type(type), l(l), r(r), idx(idx), w(w){}
};

int vCnt;
int uf_par[8'000'005];
int depth_node[2'000'005]; // 해당 깊이를 상징하는 노드 번호
vector<SweepEvent> sweepEvent[8'000'005];

int dsu_find(int x){
    if(x==uf_par[x]) return x;
    return uf_par[x] = dsu_find(uf_par[x]);
}

void makeUFTree(){
    for(int i=0; i<=8'000'000; i++) uf_par[i] = i;

    int maxDepth = *max_element(depth+1, depth+n+1);
    vCnt = maxDepth + 1;
    for(int i=1; i<=vCnt; i++){
        uf_par[i] = i;
        depth_node[i-1] = i;
    }

    vector<Chunk> chunks;
    for(int i=1; i<=n; i++){
        chunks.push_back(Chunk(baseTree.in[i], depth[i]+1, A[i]));
    }

    vector<tuple<int, int, int> > surveys;
    vector<pair<int, int> > uf_edges;
    for(int i=1; i<=q; i++){
        if(queries[i].type == 1){ // relocation
            int a = queries[i].a, b = queries[i].b; // a -> b
            if(!depth_node[a]) continue;
            int da = depth_node[a];
            assert(da == dsu_find(da));

            if(!depth_node[b]){ // b가 존재하지 않는 경우
                int db = depth_node[b] = ++vCnt;

                uf_par[da] = db;
                uf_edges.push_back({da, db});
                depth_node[a] = 0;
            }
            else{
                int db = depth_node[b];
                assert(db == dsu_find(db));

                int idx = ++vCnt;
                uf_par[da] = uf_par[db] = idx;
                uf_edges.push_back({da, idx});
                uf_edges.push_back({db, idx});
                depth_node[a] = 0, depth_node[b] = idx;
            }
        }
        else if(queries[i].type == 2){ // immigration
            int d = depth[queries[i].a];
            if(depth_node[d]){
                int idx = ++vCnt;
                int da = depth_node[d];
                assert(da == dsu_find(da));
                chunks.push_back(Chunk(baseTree.in[queries[i].a], idx, queries[i].b));

                uf_par[da] = idx;
                depth_node[d] = idx;
                uf_edges.push_back({da, idx});
            }
            else{
                int idx = ++vCnt;
                depth_node[d] = idx;
                chunks.push_back(Chunk(baseTree.in[queries[i].a], idx, queries[i].b));
            }
        }
        else{ // survey
            int d = depth[queries[i].a];
            if(!depth_node[d]) continue;

            surveys.push_back({queries[i].a, depth_node[d], i});
        }
    }

    { // ufTree 정리
        int root = ++vCnt;
        for(int i=0; i<=n; i++){
            if(depth_node[i]){
                uf_edges.push_back({depth_node[i], root});
            }
        }

        ufTree.init(vCnt, uf_edges);
        ufTree.dfsOrder(root);
    }

    { // chunk와 survey 정리. event 0이 항상 event 1보다 먼저 들어간다.
        for(auto p: chunks){
            p.in2 = ufTree.in[p.in2];
            sweepEvent[p.in2].push_back(SweepEvent(0, p.in1, p.in1, 0, p.val));
        }
        for(auto &[nodeIdx, ufIdx, queryIdx]: surveys){
            int in = ufTree.in[ufIdx], out = ufTree.out[ufIdx];
            int l = baseTree.in[nodeIdx], r = baseTree.out[nodeIdx];
            sweepEvent[in-1].push_back(SweepEvent(1, l, r, queryIdx, -1));
            sweepEvent[out].push_back(SweepEvent(1, l, r, queryIdx, 1));
        }
    }
}

struct Fenwick{
    int n;
    int tree[10'000'002];

    void init(int _n){
        n = _n;
        fill(tree, tree+n+1, 0);
    }

    void add(int x, int v){
        while(x<=n){
            tree[x] += v;
            x += x&-x;
        }
    }

    int sum(int x){
        int ret = 0;
        while(x){
            ret += tree[x];
            x -= x&-x;
        }
        return ret;
    }

    int sum(int l, int r){
        return sum(r) - sum(l-1);
    }
} tree;

int ans[2'000'002];

void sweep(){
    tree.init(n);

    for(int y=0; y<=vCnt; y++){
        for(SweepEvent &ev: sweepEvent[y]){
            if(ev.type == 0) tree.add(ev.l, ev.w);
            else ans[ev.idx] += tree.sum(ev.l, ev.r) * ev.w;
        }
    }

    for(int i=1; i<=q; i++){
        if(queries[i].type == 3) cout << ans[i] << '\n';
    }
}