重み付きUnionFind(ポテンシャル付きUnionFind)は差分制約系 (牛ゲー) の処理が出来る。
Union(x, y, 10); Union(y, z, 20); を実行すると Diff(x, z) の戻り値で30が返ってくる。
この性質を利用するとABC087Dなどがライブラリをコピペするだけで解ける。
https://atcoder.jp/contests/abc087/submissions/10695617
最初は中身が全く理解できなかったので関数の仕様だけ暗記しようかと思っていたが、20時間ほどかけて色々と考察したら何となく分かったような気がしてきたのでメモを書いておく。
なお、自分が納得できれば十分だったので内容の正しさは一切保証しない。
まず、1本の木における動作についての考察。例えば、頂点1から4まで接続された木構造にUnion(3,5,1)をすると頂点3の重みに+1したのが頂点5の重みとなる。Union(4,6,9)でも同様の処理が出来る。よく見ると累積和で矛盾なく成立している。すごい。
次に、2本の木における動作についての考察。Union(x,y,10)を考えてみる。UnionFindの性質上xとyを併合すると言っても実際に接続されるのはxの根とyの根である。したがって、yの木サイズが小さい場合、Union内で更新しなければならない値はxの根と接続されるyの根の重みになる。yの木サイズが小さい場合の方程式はnew_w = 100 + w - 50 = 60になる。木サイズが逆なら符号が逆になる。
この60は意味不明だと思うが、この答えnew_wの意味はDiffを呼び出したときに分かる(後述)。また「他の頂点の重みを更新しないとDiff(x,y)が正しく得られないのでは」という悩みもDiffの動作説明で解決する。
T Diff(long long x, long long y){ return Weight(y) - Weight(x); }
T Weight(long long v){ Find(v); return weight[v]; }
long long Find(long long v)
{
if (root[v] == v) return v;
long long r = Find(root[v]);
weight[v] += weight[root[v]];
return root[v] = r;
}
Diff(x,y)を呼び出すと内部でWeight関数が呼ばれ、Weight内ではFind関数が呼ばれる。Findは根が得られるまで再帰的に呼ばれるが、UnionFindは木の高さが低くなりやすい性質があるので、一般的な構造で図よりFind(y)が1回だけ呼ばれたとする。そのときFind内では weight[y] += weight[root[y]] により、yの重みにyの根の重み(前述のnew_w)が加算される。ここで、ついに頂点yの重みがxの根から累積的に正しい値に修正される。したがって、Diff(x,y)が正しい値を返してくれる。なんだこれ神か。