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

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

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

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

問題概要

 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 の各要素に対応する
  •  i (=0, 1, \dots, M-1) 回目の併合において、要素  a, b を併合するとして、頂点  i + N を、この併合操作を表す頂点とする
    • まず、 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;
}