読者です 読者をやめる 読者になる 読者になる

AGC002 Stamp Rally

http://agc002.contest.atcoder.jp/tasks/agc002_d
 本当はこんなことしている場合ではないのですが、息抜きにこの前のAtCoderの問題を解いたので久しぶりに更新。
 問題はN 頂点 M 辺の無向グラフ(辺は1~Mの番号が振られている)が与えられたとき、以下のクエリに答えよ、という問題。
・頂点x, yから辺を辿ってz個以上の頂点を巡るとき、通らなければいけない辺の番号の最小値。

 この問題、制約が100000と大きいので、全域木の逐次構築でも無理、にぶたんでも無理(上手にやればできるかも)、と本番中は手も足も出ませんでした。
 解説(http://agc002.contest.atcoder.jp/data/agc/002/editorial.pdf)によると並列に二分探索するとできるらしいです(二分探索する対象は求める最小値で、逐次構築する時間があればUF木で判定できる)。並列に二分探索と言われてもピンと来ないので、解説の英語のほうにのっている解法をそのまま実装しました。

#include <bits/stdc++.h>
using namespace std;
#define REP(i,s,n) for(int i=(int)(s);i<(int)(n);++i)
#define rep(i,n) REP(i, 0, n)

struct UnionFind {
    vector<int> data; int unitenum;
    UnionFind() {}
    void init(int N) {
        data.assign(N, -1);
        unitenum = 0;
    }
    void unite(int x, int y) {
        ++unitenum;
        x = root(x); y = root(y);
        if(x != y){
            data[x] += data[y];
            data[y] = x;
        }
    }
    int root(int x) { return data[x] < 0 ? x : data[x] = root(data[x]);}
    bool same(int a, int b) { return root(a) == root(b);}
    int group(int a) { return data[root(a)] * -1;}
};

typedef pair<int, int> pint;

int get(UnionFind& uf, int a, int b) {
    if(uf.same(a, b)) 
      return uf.group(a);
    else
      return uf.group(a) + uf.group(b);
}

struct query {
    int x, y, z, num;
    query(){}
    query(int x, int y, int z, int n)
         : x(x), y(y), z(z), num(n) {}
};

int main() {
    cin.tie(0); ios::sync_with_stdio(false);
    int N, M; cin >> N >> M;
    vector<pint> G(M);
    rep(i, M) {
        int a, b; cin >> a >> b;
        --a; --b;
        G[i] = pint(a, b);
    }
    int Q; cin >> Q;
    vector<query> qs(Q);
    rep(i, Q) {
        int x, y, z;
        cin >> x >> y >> z;
        --x; --y;
        qs[i] = query(x, y, z, i);
    }
    
    queue<tuple<int, int, vector<query>>> que;
    que.emplace(0, M, qs);
    vector<int> ans(Q, 0);
    UnionFind uf;
    uf.init(N);
    while(!que.empty()) {
        int L, R; vector<query> cur;
        tie(L, R, cur) = que.front(); que.pop();
        if(R - L == 1) {
            for(const auto& q : cur)
              ans[q.num] = L;
            continue;
        }
        int mid = (R + L) >> 1;
        if(uf.unitenum > mid) uf.init(N);
        while(uf.unitenum < mid) {
            int a, b;
            tie(a, b) = G[uf.unitenum];
            uf.unite(a, b);
        }
        vector<query> Lbox, Rbox;
        for(const auto& q : cur) {
            int score = get(uf, q.x, q.y);
            if(score < q.z)
              Rbox.push_back(q);
            else
              Lbox.push_back(q);
        }
        que.emplace(L, mid, Lbox);
        que.emplace(mid, R, Rbox);
    }
    for(int k : ans) cout << k + 1 << '\n';
    return 0;
}

 まずキューに(L, R, クエリが入ったvector)をつっこんでおきます。vector中のクエリ全部に対し、求める最小値が半開区間(L, R]にあるようにします。このとき、uf木を番号(L + R) / 2の辺まで構築すれば、vector中のクエリを仕分けできます。そして、幅優先的に小さい方からやっていくので(別にDFSでも変わりませんが)UF木が構築される回数はたかだかlogM回、と(というかL, Rについて二分探索なので当然ですね)。なんて頭のいい解法なんだ……。個人的に、かなり面白いと思いました。