Mister雑記

競プロやります。

AGC002-D 「Stamp Rally」 (1000)

おや?Union-Findの様子が......

問題リンク

概要

 N頂点 M辺の連結グラフがある。辺には 1 \sim Mの独立した番号がついており、 i番目の辺は頂点 a_i b_iを結んでいる。

このとき以下のようなクエリが Q個与えられる。これらを全て処理せよ。

  •  j番目のクエリは3つの整数 x_j, y_j, z_jからなる。
  •  x_j, y_jからそれぞれ辺を辿って、合計 z_j個の頂点を訪れることを考える。
    • 始点が別でも同じ頂点は1つと数える。
  • このとき辿った辺の番号の最大値がスコアとなる。
  • スコアを最小化したとき、そのスコアを求めよ。

制約

  •  3 \leq N \leq 10^5
  •  N - 1 \leq M \leq 10^5
  •  1 \leq Q \leq 10^5
  •  3 \leq z_j \leq N

解説

editorialと違ってライブラリで殴る感じの解法です。ライブラリ不要な解法はeditorialを参照。

可能性問題への帰着

問題のように最小化をそのまま考えるのは難しいので、 x, y, zが与えられたときに「スコア s以下で実現可能か?」という問題を考えてみる。 これを解くのは割と簡単で、辺 1 \sim sのみからなるグラフについて x, yを含む連結成分の個数を数え、これと zを比較すればいい。

連結成分を数えるのはやはりUnion-Find(以降「UF」)を使うのが一番楽だろう。ということで、

  • UFを初期化する
  • 辺を 1 \sim sまで追加する
  •  x, yが連結しているか調べ、
    • 連結しているなら xを含むグループのサイズ
    • 連結していないなら x, y を含むグループのサイズの和
  •  z以上ならOK、そうでなければNG

とすればこの問題は解ける。

C++擬似コードを書くなら以下のような感じ。

bool judge (int x, int y, int z, int s) {
    
    UnionFind uf(N + 1); // union-findを頂点数Nで初期化
    for (int i = 1; i <= s; ++i) {
        uf.unite(a[i], b[i]);
    }
        
    int con = uf.size(x); // 連結成分の個数
    if (uf.find(x) != uf.find(y)) con += uf.size(y);
        
    return con >= z;
}

さらにスコアには単調性がある、すなわち「スコア sで実現可能なら任意の s' \geq sでも実現可能」という性質が明らかに成り立つので、スコアを二分探索することで高速に最小値が求められる。

部分永続union-find

これで各クエリをUFの構築等に O(M \log N)、解の二分探索に O(\log M)で計 O(M \log M \log N)で解けるようになった。 しかし Qが大きいためこれでは全然間に合わない。

上の解法で明らかに無駄なのは、同じUFをクエリを跨いで何度も構築しているところにある。 そこで各 1 \leq s \leq Mについて辺 1 \sim sからなるUFを全部メモしておけば、可能性判定を O(\log M)でできるようになって十分高速になる。

愚直に配列でUFを持つと空間量が O(MN)になって明らかにMLEなのだが、実はこれを通常のUFと同じ O(N)の空間量で実現できるデータ構造が存在する。それが部分永続union-find(以降「部分永続UF」)である。

部分永続UFは具体的には以下のクエリに対処できる:

  •  unite(x, y)
    • 頂点 x, yを結合する
  •  find(x, t)
    •  t uniteを行った時点での、 xの親を探す
    • この際経路圧縮は行われない
  •  size(x, t)
    •  t回uniteを行った時点での、 xを含むグループの頂点数を返す

上にも書いた通り、部分永続UFではその性質上経路圧縮が行えないことに注意が必要である。unite時にrankを意識しないと偏った入力が与えられたときにTLEになってしまう。

部分永続UFの解説は長くなるので、実装だけ別の記事にまとめた。なんにせよこれが構築できればあとは少し手を加えるだけでACとなる。

解答例

#include <vector>
#include <iostream>
using namespace std;

using ll = long long;

#define SIZE(v) (static_cast<int>((v).size()))
#define ALL(v) (v).begin(), (v).end()
const int INF = 1 << 25;

// 部分永続UFのオブジェクト
class persistentUF {
public:
    explicit persistentUF(int N) : V_NUM(N), now(0) {
        par.resize(N);
        for(int i = 0; i < V_NUM; ++i) {
            par[i] = i;
        }
        
        rank.resize(N);
        fill(ALL(rank), 0);

        time.resize(N);
        // 親がいないときはtimeはINFにしておく
        fill(ALL(time), INF);

        num.resize(N);
        for (int i = 0; i < V_NUM; ++i) {
            // 時刻0にて各要素の要素数は1
            num[i].push_back(make_pair(0, 1));
        }
    }

    // 時刻tにおけるxの親を返す
    ll find(int x, int t) {
        if (time[x] > t) return x;
        return find(par[x], t);
    }

    // 時刻tにおけるxを含むグループの要素数を返す
    ll size(int x, int t) {
        x = find(x, t);
        ll ok = 0, ng = SIZE(num[x]);

        // 二分探索でサイズが入ったindexを探す
        while (ng - ok > 1) {
            ll mid = (ok + ng) / 2;
            if (num[x][mid].first <= t) {
                ok = mid;
            } else {
                ng = mid;
            }
        }
        return num[x][ok].second;
    }

    // xとyを含むグループを統合する
    void unite(int x, int y) {
        ++now;
        x = find(x, now);
        y = find(y, now);

        if (x == y) return;

        // rankの大きい方にくっつける
        if (rank[x] > rank[y]) swap(x, y);

        // numを更新する前に親を更新するとバグるので注意
        num[y].push\_back(make\_pair(now, size(x, now) + size(y, now)));

        par[x] = y;
        time[x] = now;
        if (rank[x] == rank[y]) ++rank[y];
    }

    // xとyが時刻tにおいて同じグループに属するか判定
    bool same(int x, int y, int t) {
        return find(x, t) == find(y, t);
    }

    int V_NUM, now;
    // 頂点数、時刻もとい現在追加した辺の数
    
    vector<int> par, rank, time;
    // 直接の親、子の深さ、親が更新された時刻
    
    vector<vector<pair<int, int>>> num;
    // 時刻とグループの頂点数のペアを格納
};


int main() {
    int N, M;
    cin >> N >> M;

    // 部分永続UFの初期化
    persistentUF uf(N + 1);
    for(int i = 0; i < M; ++i) {
        int a, b;
        cin >> a >> b;
        uf.unite(a, b);
    }

    int Q;
    cin >> Q;
    for(int i = 0; i < Q; ++i) {
        int x, y, z;
        cin >> x >> y >> z;
        
        // 二分探索で最小値を求める
        int ok = M, ng = 0;
        while (ok - ng > 1) {
            // 上に書いた擬似コードど同じ
            int mid = (ok + ng) / 2;
            int con = uf.size(x, mid);;
            if (!uf.same(x, y, mid)) con += uf.size(y, mid);

            if (con >= z) {
                ok = mid;
            } else {
                ng = mid;
            }
        }
        cout << ok << endl;
    }
    return 0;
}

感想

  • 部分永続UF強すぎる。
  • ずっと部分永続UFを使わずに済む方法を考えていただけに、editorialの解法にびっくり。