Weighted Union-Find

Info

グラフの連結性と同集合内の2頂点のポテンシャルの差を判定できるデータ構造

Complexity

辺の追加 - \(O(\alpha(n))\) (\(\alpha はアッカーマン関数の逆関数\))
頂点が所属してる集合の名前の取得 - \(O(\alpha(n))\)
頂点が所属してる集合の大きさの取得 - \(O(\alpha(n))\)
同集合内の2頂点のポテンシャルの差の取得 - \(O(\alpha(n))\)

Interface

WeightedUnionFind(int n) 大きさ\(n\)の頂点集合で初期化する
find(int i) 頂点\(i\)の所属してる集合の名前を返す
unite(int a, int b, long long w) \(a\)-\(b\)にポテンシャル差\(w\)の辺を張る
diff(int a, int b) \(a\)-\(b\)のポテンシャル差を返す
size[find(int i)] 頂点\(i\)の所属してる集合の大きさを返す

Code

using lnt = long long;
/* <WeightedUnionFind> */

struct WeightedUnionFind
{
	std::vector<int> par,size;
	std::vector<lnt> weight;
	WeightedUnionFind(int n) : par(n),size(n,1),weight(n,0) {
		std::iota(par.begin(),par.end(),0);
	}
	int find(int i) {
		if(par[i]==i) return i;
		int x=find(par[i]);
		weight[i]+=weight[par[i]];
		return par[i]=x;
	}
	int unite(int i, int j, lnt w) {
		if(same(i,j)) {
			if(diff(i,j)!=w) {
				return 1;
			} return 0;
		}
		int x=find(i), y=find(j);
		if(size[x]<size[y]) {
			std::swap(x,y);
			std::swap(i,j);
			w=-w;
		}
		par[y]=x;
		weight[y]=w+weight[i]-weight[j];
		return 0;
	}
	lnt diff(int i, int j) {
		find(i); find(j);
		return weight[j]-weight[i];
	}
	bool same(int i, int j) {
		return find(i)==find(j);
	}
};

/* <WeightedUnionFind> */

Examples

WeightedUnionFind uf(10);
uf.unite(0,9,1);
uf.unite(1,2,1);
uf.unite(1,3,5);
uf.unite(4,2,-4);
std::cout << uf.diff(3,4) << '\n'; // 0
std::cout << uf.diff(1,4) << '\n'; // 5