マージテクを知らなくても雰囲気で通せるのかもしれない

問題へのリンク

問題概要

$N$ 人の生徒がそれぞれクラス $C_{1}, C_{2}, \dots, C_{N}$ に所属している。

各生徒はそれぞれの家から出発したあと、他の生徒と合流を繰り返しながら学校へ向かう。一度合流した生徒が分かれることはない。

$Q$ 個のクエリが与えられるので、順番に処理せよ。クエリには以下の 2 種類がある。

1 $a$ $b$ ：生徒 $a$ を含む集団と、生徒 $b$ を含む集団が合流する (既に合流しているときは何も起こらない)
2 $x$ $y$ : クエリの時点で既に生徒 $x$ と合流している生徒(生徒 $x$ を含む)のうち、クラス $y$ に属している生徒の数を求める

制約

$1 \le N, Q \le 2 \times 10^{5}$

解法 (1)：マージテク

マージテクとは「データ構造を併合していくときに、併合の仕方を工夫することで計算量を小さく抑える」というテクニック。しかしながら、それを知らなくても自然に DP したら通った、という人も多かったかもしれない！

とりあえずこの手の問題は、まずは計算量が多少悪くても、とりあえず解ける方法を考えるのがよいと思う。そして、クエリの各過程において、次のデータを管理していけば解けそう。

dp[ a ][ c ] := 生徒 a が所属しているグループに、クラス c の人は何人いるか？

実際には、このデータは生徒全員に持たせる必要はない。Union-Find では各グループを「根付き木」で管理するので、根にだけ持たせれば OK。そうすれば、生徒 a が所属しているグループに関する情報は次のようにして取得できる。

dp[ uf.root(a) ][ c ]

さて、クエリ 1 の併合 (生徒 a のいるグループと、生徒 b のいるグループのマージ) によって、データがどのように移り変わるのかを考えよう。まず、生徒 a, b が、それぞれのグループの「根」で置き換えておくことにする。そして、次のようにすると仮定する。

生徒 a, b のうち、生徒 a が「根」になるようにマージすることにする
- つまり、生徒 a が生徒 b の親になる

このとき、次のように更新すれば OK。

各クラス c について、

dp[ a ][ c ] += dp[ b ][ c ]

実際には、c は全クラスを動く必要はなく、生徒 b のいるグループのメンバーの属するクラスのみ動かせば OK。よって、dp[ a ] や dp[ b ] は map 型でもつとよい。このとき、dp[ a ][ c ] += dp[ b ][ c ] という更新は、あたかも「集合の併合」のような挙動になる。たとえば

dp[ a ] = {0: 3, 1: 5, 2: 2}
dp[ b ] = {1: 4, 3: 1}

という感じだったら、更新後には

dp[ a ] = {0: 3, 1: 9, 2: 2, 3: 1}

となる。これは集合 {1, 3} を集合 {0, 1, 2} にマージしたような感じになっている。このようにデータ構造を併合していくときに一般的に使えるテク (通称、マージテク) がある。実は Union-Find の併合規則では、すでにそのテクニックが使われている。よって今回は、Union-Find の併合規則の挙動の通りに併合するのでも OK。つまり、

Union-Find において、根 b が根 a の子供となるように (根 a が新たな根となるように) 併合する、となったとき
dp[ b ] を dp[ a ] へと併合する

というふうにすれば OK。

たとえば極端な場合として、 $N$ 人全員が別々のクラスに属する場合が最悪となる。生徒 i が属するクラスを i としても一般性を失わない。このとき、辞書 dp の併合過程と、Union-Find 上の集合の併合過程とは、ピッタリ対応することになる。そして Union-Find の併合規則に従うとき、任意の要素 i に対して、それがサイズの小さい map 側 (dp[ b ] 側) から、サイズの大きい map 側 (dp[ a ] 側) へと併合される回数は、 $O(\log N)$ で抑えられることが言えるのだ。

以上から、全体の計算量も $O(N (\log N)^{2} + Q \log N)$ で抑えられることがわかった。 $O((\log N)^{2})$ というふうに二乗が付くのは、dp テーブルを map で管理しているため。

コード

#include <bits/stdc++.h>
using namespace std;

struct UnionFind {
    vector<int> par;

    UnionFind() { }
    UnionFind(int n) : par(n, -1) { }
    void init(int n) { par.assign(n, -1); }
    
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y) {
        x = root(x); y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y); // merge technique
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }
};

int main() {
    int N, Q;
    cin >> N >> Q;
    vector<int> C(N);
    for (int i = 0; i < N; ++i) cin >> C[i], --C[i];

    UnionFind uf(N);
    vector<map<int, int>> dp(N);
    for (int i = 0; i < N; ++i) dp[i][C[i]] += 1;
    for (int _ = 0; _ < Q; ++_) {
        int type, x, y;
        cin >> type >> x >> y;
        --x, --y;
        if (type == 1) {
            if (uf.issame(x, y)) continue;

            // とりあえず根に移動して、マージする
            x = uf.root(x), y = uf.root(y);
            uf.merge(x, y);

            // x が根になるようにする
            if (uf.root(y) == y) swap(x, y);

            // y の値を x にマージ
            for (auto it : dp[y]) dp[x][it.first] += it.second;
        }
        else {
            // 生徒 x、クラス y
            int r = uf.root(x);
            if (dp[r].count(y)) cout << dp[r][y] << endl;
            else cout << 0 << endl;
        }
    }
}

解法 (2)：Union-Find のマージ過程を表す木

Union-Find の要素数を $N$ 、併合回数を $M$ 回としたとき、以下のような頂点数 $N+M$ の根付き森を考えよう。

頂点 $0, 1, \dots, N-1$ については、Union-Find の各要素に対応する
回目の併合において、要素を併合するとして、頂点を、この併合操作を表す頂点とする
- まず、 $a, b$ をそれぞれその時点での Union-Find の根に置き換える
- さらに、頂点 $a$ が根になった最後の瞬間の併合を表す頂点を $c$ とし、頂点 $d$ が根になった最後の瞬間の併合を表す頂点を $d$ とする
- 頂点 $i + N$ の子頂点を頂点 $c, d$ とする

最終的には「最後の併合を表す頂点」を根とする、頂点数 $N + M$ の根付き森となる。追加で必要なデータとして

ancestor[ v ] := 頂点 $v$ ( $= 0, 1, \dots, N-1$ ) が根となるような併合を行った最後の瞬間の、併合操作を表す頂点番号 ( $N, N+1, \dots, N+M-1$ のいずれかとなる)

をもつ。こうしてできる根付き森は、Union-Find の各併合過程の状態を保存したものとなっていることから、さまざまなことができる。強いのは、

「Union-Find の併合過程を表す木 (森) 上の Euler Tour をとることによって、各併合過程のその時点におけるグループを『区間』として扱える」

という点だ。今回も、各併合過程におけるグループのクラスの人数は、「区間の総和」として表されるわけだ。よって、Euler Tour 上の累積和を考えてあげれば、その差分によって求められる。

ただし今回は更新クエリはない (途中でクラス変更するなどがない) ため、セグメント木を持ち出す必要まではない。次のようにすれば OK。ここで、res[ i ] を i 番目のクエリに対する答えとしている (初期値を 0 とする)。

Euler Tour 上で常にその時点でたどった頂点についての、各クラスの人数の累積和を表す vector<int> 型変数 counter を管理しておく
- 要素数はクラス数
親頂点 $p$ から子頂点 $v$ ( $\ge N$ ) へと入る瞬間 (行きがけ順) には、頂点 $v$ に関するクラス $c$ についてのクエリに対して、res[ (クエリ番号) ] -= counter[c] としておく
頂点 $v$ から親頂点へと戻る瞬間 (帰りがけ順) には、頂点 $v$ に関するクラス $c$ についてのクエリに対して、res[ (クエリ番号) ] += counter[c] としておく

これを 1 回回せば、すべてのクエリに対する答えが自動的に求められる。計算量は $O(N + Q\alpha(N))$ 。

コード

#include <bits/stdc++.h>
using namespace std;

struct UnionFind {
    vector<int> par;

    UnionFind() { }
    UnionFind(int n) : par(n, -1) { }
    void init(int n) { par.assign(n, -1); }
    
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y, int id = -1) {
        x = root(x); y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y); // merge technique
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }
};

using Query = pair<int,int>; // id, class
int main() {
    int N, Q;
    cin >> N >> Q;
    vector<int> C(N);
    for (int i = 0; i < N; ++i) cin >> C[i], --C[i];

    // build forest
    UnionFind uf(N);
    vector<int> left(N+Q, -1), right(N+Q, -1), ancestor(N+Q);
    iota(ancestor.begin(), ancestor.end(), 0);
    vector<vector<Query>> queries(N+Q);
    int mergeNum = N, queryNum = 0;
    for (int q = 0; q < Q; ++q) {
        int type, x, y;
        cin >> type >> x >> y;
        --x, --y;
        if (type == 1) {
            x = uf.root(x), y = uf.root(y);
            if (x == y) continue;
            uf.merge(x, y);

            // x が親となるように
            if (uf.root(x) != x) swap(x, y);

            left[mergeNum] = ancestor[x];
            right[mergeNum] = ancestor[y];
            ancestor[x] = mergeNum++;
        } else {
            x = uf.root(x);
            int v = ancestor[x];
            queries[v].emplace_back(queryNum++, y);
        }
    }
    
    // dfs
    vector<int> res(queryNum, 0), counter(N, 0);
    auto dfs = [&](auto self, int v) -> void {
        for (auto q : queries[v]) res[q.first] -= counter[q.second];
        if (v < N) counter[C[v]]++; // leaf
        else self(self, left[v]), self(self, right[v]);
        for (auto q : queries[v]) res[q.first] += counter[q.second];
    };
    for (int v = 0; v < N; ++v) {
        if (uf.root(v) != v) continue;
        dfs(dfs, ancestor[v]);
    }
    for (auto val : res) cout << val << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 183 F - Confluence (水色, 600 点)

問題概要

制約

解法 (1)：マージテク

コード

解法 (2)：Union-Find のマージ過程を表す木

コード