mayoko’s diary

プロコンとかいろいろ。

AtCoder Grand Contest 002 D - Stamp Rally

解法

クエリが 1 個の場合を考えると, UnionFind を使うことで O((N+M)logM) ぐらいで答えを求めることができます(単体だけを考えたら O(N+M) でできそうだけどそれは今はいいです)。

というのは, ok(m) = (辺の id が m の部分までグラフに辺を付け加えたとき, x, y の連結成分の合計が z を超えるか?)というのを m について二分探索すれば良いですね。ただ, これを各クエリにやっていると間に合いません。

そこで, 上記の方法で二分探索を並列にやっていくことを考えます(なんかパラレルバイナリサーチとか名前ついてるっぽい?)。

各クエリに対して, 「今解は (l, r] に絞られている」というのを覚えておきます。すると, 次に調べたい中間地点は m = (l+r)/2 です。この m を各クエリで覚えておきます。

すると, 一度「辺の id を 0 から順番に増やしていきながら頂点を連結させていく」という作業をやるたびに, 各クエリの ok(m) を判定することができます(普通に辺を張っていき, 調べるべき m が来たら判定する感じ)。よって, O((N+M+Q)logM) くらいですべてのクエリの判定ができるわけです。

struct UnionFind {
    vector<int> par;
    int n, cnt;
    UnionFind(const int& x = 0) {init(x);}
    void init(const int& x) {par.assign(cnt=n=x, -1);}
    inline int find(const int& x) {return par[x] < 0 ? x : par[x] = find(par[x]);}
    inline bool same(const int& x, const int& y) {return find(x) == find(y);}
    inline bool unite(int x, int y) {
        if ((x = find(x)) == (y = find(y))) return false;
        --cnt;
        if (par[x] > par[y]) swap(x, y);
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    inline int count() const {return cnt;}
    inline int count(int x) {return -par[find(x)];}
    int count(int x, int y) {
        if (same(x, y)) return count(x);
        return count(x) + count(y);
    }
};

struct Query{
    int x, y, z;
    int l, r;
};

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int N, M;
    cin >> N >> M;
    vector<pii> es(M);
    for (int i = 0; i < M; i++) {
        cin >> es[i].first >> es[i].second;
        es[i].first--; es[i].second--;
    }
    int Q;
    cin >> Q;
    vector<Query> qs(Q);
    for (int i = 0; i < Q; i++) {
        cin >> qs[i].x >> qs[i].y >> qs[i].z;
        qs[i].x--; qs[i].y--;
        qs[i].l = -1, qs[i].r = M-1;
    }
    for (int t = 0; t < 30; t++) {
        vector<vi> check(M);
        for (int i = 0; i < Q; i++) {
            int m = (qs[i].l + qs[i].r)/2;
            check[m].push_back(i);
        }
        UnionFind uf(N);
        for (int i = 0; i < M; i++) {
            uf.unite(es[i].first, es[i].second);
            for (int j = 0; j < check[i].size(); j++) {
                int id = check[i][j];
                int cnt = uf.count(qs[id].x, qs[id].y);
                if (cnt >= qs[id].z) {
                    qs[id].r = i;
                } else {
                    qs[id].l = i;
                }
            }
        }
    }
    for (int i = 0; i < Q; i++) {
        cout << qs[i].r+1 << endl;
    }
    return 0;
}