AtCoder Grand Contest 002 D - Stamp Rally

問題

D: Stamp Rally - AtCoder Grand Contest 002 | AtCoder
N頂点，M辺からなるグラフがあり，Q組の兄弟がスタンプラリーを行う．i番目の兄弟はそれぞれXi，Yiにいて合わせてZiのスタンプを集めたいと思っている．通ることになる辺の最大の番号を最小にする問題

解法

まず，クエリ１つずつに対して最小値を求めることを考える．Union-Findを使って，頂点の個数の合計がZiを超えるまで順番に辺を加えていけばいい．しかし，この解法の計算量はO(MQ)になってしまう．そこでクエリをまとめて計算をすることを考える．Union-Findで辺をM/2まで追加し，その時点でそれぞれのクエリがZiを超えているか判定し，超えたものの集合と超えていない集合に分ける．これを図のようにlogM回繰り返すと，最終的にj番目まで辺を加えたときにZiを超えるクエリの集合を求められる．図でいうと行ごとにO(M+Q)かかるので全体でO((M+Q)logM)に抑えられる．
コードを書く際はMを2のべき乗まで広げると実装しやすくなる．

永続Union-Findを使っても解けるらしい

ソースコード

#include <iostream>
#include <vector>
#include <map>
#include <set>
#include <queue>
#include <algorithm>
#include <iomanip>
#include <cassert>
using namespace std;

#define GET_ARG(a,b,c,F,...) F
#define REP3(i,s,e) for (i = s; i <= e; i++)
#define REP2(i,n) REP3 (i,0,(int)(n)-1)
#define REP(...) GET_ARG (__VA_ARGS__,REP3,REP2) (__VA_ARGS__)
#define RREP3(i,s,e) for (i = s; i >= e; i--)
#define RREP2(i,n) RREP3 (i,(int)(n)-1,0)
#define RREP(...) GET_ARG (__VA_ARGS__,RREP3,RREP2) (__VA_ARGS__)
#define DEBUG(x) cerr << #x ": " << x << endl

typedef long long ll;

constexpr int INF = 1e8;
constexpr int MOD = 1e9+7;

int par[100000];
int cnt[100000];
vector<int> bs[262143];
int A[100000], B[100000];
int X[100000], Y[100000], Z[100000];
int ans[100000];

void init(int n) {
    int i;
    REP (i,n) par[i] = i;
    REP (i,n) cnt[i] = 1;
}

int find(int x) {
    if (par[x] == x) return x;
    else return par[x] = find(par[x]);
}

void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x != y) cnt[y] += cnt[x];
    par[x] = y;
}

int count(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y) return cnt[x];
    else return cnt[x] + cnt[y];
}

int main(void) {
    int i, N, M;
    scanf("%d%d",&N,&M);
    REP (i,M) {
        scanf("%d%d",&A[i],&B[i]);
        A[i]--; B[i]--;
    }
    int Q;
    scanf("%d",&Q);
    REP (i,Q) {
        scanf("%d%d%d",&X[i],&Y[i],&Z[i]);
        X[i]--; Y[i]--;
    }

    REP (i,Q) bs[0].push_back(i);
    int size = 1;
    while (M > size) size <<= 1;
    for (int j = 0, k = 0; 1 << k < size; k++) {
        init(N);
        int s = 0, e = size / (1<<k);
        REP (i,size) {
            if (i == (s + e) / 2) {
                for (auto a: bs[j]) {
                    if (count(X[a],Y[a]) >= Z[a])
                        bs[j*2+1].push_back(a);
                    else
                        bs[j*2+2].push_back(a);
                }
                j++;
                s = e; e += size / (1<<k);
            }
            if (i < M) unite(A[i],B[i]);
        }
    }
    REP (i,M) for (auto x: bs[i+size-1]) {
        ans[x] = i+1;
    }
    REP (i,Q) printf("%d\n",ans[i]);
    return 0;
}