このデータ構造の実装には、 けんちょんさんの解説 を大変参考にさせていただきました。
/*
この重み付きunion_find木で出来ること
例)weight_union_find tree(N);などで宣言(N:頂点数)
tree.same(x,y):要素xと要素yが同じグループに属するか判定
tree.unite(x,y,w):weight(y)=weight(x)+wとなるように木を併合
tree.diff(x,y):weight(y)-weight(x)を返す
tree.size(x):要素xが属する木の頂点数を返す
*/
struct weight_union_find
{
vector<long long> par; //親の番号
vector<long long> rank; //木の深さ(根のランクは0)
vector<long long> siz; //要素xが根なら木の頂点数を格納する
vector<long long> par_diff; //親と子のコストの差
//初期化子リストを用いた初期化
weight_union_find(long long N) : par(N), rank(N), siz(N), par_diff(N)
{
for (long long i = 0; i < N; i++)
{
par[i] = i;
rank[i] = 0;
siz[i] = 1;
par_diff[i] = 0;
}
}
//要素xが所属する木の根を再帰的に発見する
long long root(long long x)
{
if (par[x] == x)
return x;
else
{
//経路圧縮+累積和
long long ret = root(par[x]);
par_diff[x] += par_diff[par[x]];
return par[x] = ret;
}
}
//経路圧縮+コストを返す
long long weight(long long x)
{
root(x);
return par_diff[x];
}
//yのコスト-xのコスト
long long diff(long long x, long long y)
{
return weight(y) - weight(x);
}
//weight(y)=weight(x)+wとなるように木を併合
void unite(long long x, long long y, long long w)
{
long long rx = root(x);
long long ry = root(y);
w += weight(x);
w -= weight(y);
if (rx == ry)
return; //同じ木に属してたらそのまま
if (rank[rx] < rank[ry])
{
swap(rx, ry);
siz[ry] = siz[rx] + siz[ry];
w = -w; //符号反転
}
else
{
siz[rx] = siz[rx] + siz[ry];
}
if (rank[rx] == rank[ry])
rank[rx]++;
par[ry] = rx;
par_diff[ry] = w;
}
//要素xが属する木と要素yが属する木が同じならtrueを返す
bool same(long long x, long long y)
{
return root(x) == root(y);
}
//要素xが属する木の頂点数を返す
long long size(long long x)
{
return siz[root(x)];
}
};