平方分割を使ってAGC-DをC++で殴る

AGC002-D Stamp Rallyで平方分割入門をしようね
想定解の難しい問題を平方分割パンチC++キックでやっつけたい人用の記事
当然ネタバレになるのでいいよって人だけ読んでね

問題概要

N頂点の辺がないグラフに辺をM本順番に張っていく
各 $j(1\leq j \leq Q)$ について、 $x_j$ か $y_j$ と連結な頂点が $z_j$ 個以上に初めてなるのは辺を何本張った時か
制約: $N,M,Q \leq 10^5$

愚直大好きクラブ

連結成分を見るのでUnionFindを使う
M本辺を貼り、一本張るごとに各 $j$ について確認する
計算量は（UnionFindはO(1)で動くとして） $O(MQ)$

bool OK(int j){
  return x[j]かy[j]と連結な頂点がz[j]個以上？
}
REP(i,m){
  uf.merge(a[i],b[i]);
  REP(j,q){
    if(ans[j])continue;//既に値が入っている
    if(OK(j))ans[j]=i+1;//jについて初めて条件を満たした
  }
}

雑なイメージコードだとこんな感じ

smartなAlgorithm

さっきのやり方だと一本辺を張るごとに各 $j$ について確認していたので遅かった
そこでA本辺を張るごとに確認することにする

つまり、
1. $A$ 本辺を張るごとに各 $j$ について条件を初めて満たしたか確認する
2.今回の $A$ 本で初めて条件満たした各 $j$ について、 $A$ 本のどこで初めて条件を満たしたかを確認する
としてやる

1の操作は $M/A$ 回行われるので全部で $O(QM/A)$
2の操作は各 $j$ についてたかだか $A$ 回の確認をするので $O(QA)$
よって全体で $O(QM/A+QA)=O(Q(M/A+A))$ になる
$M/A+A$ は $A=\sqrt M$ の時が一番小さそうですね
Mの平方根サイズずつで問題を分割しているのでこれを平方分割と呼びます
この時計算量は $O(Q\sqrt M)$ なので（早い言語なら）間に合いそう

実装

入力が最大のケースだけ考えればいいので最初から $A=320$ とかにしちゃっていいです
$A$ 本辺を張ってからもう一度 $A$ 本張る前の状態に戻って $A$ 本張ってあげる必要があるためUnionFindを二つ持ちます

const int A=320;

for(int t=0;t<m;t+=A){//A本の辺を張る
  for(int i=0;i<A&&t+i<m;i++)uf1.merge(a[t+i],b[t+i]);
  vector<int> v;//今回のA本で初めて条件を満たすjの集合
  REP(j,q){
    if(ans[j])continue;
    if(OK(j,uf1))v.push_back(j);
  }
  for(int i=0;i<A&&t+i<m;i++){//vに入ってるjが（今回追加したA本のうち）初めて条件を満たしたタイミングを見る
    uf2.merge(a[t+i],b[t+i]);//一本ずつ追加していく
    for(int p:v){
      if(ans[p])continue;
      if(OK(p,uf2))ans[p]=t+i+1;
    }
  }
}