重み付きUnion-Find
Union-Findに重みをつけたしたもの
けんちょんさんの記事
ABC373Dの解説
LibraryCheckerへの提出
例題
F - Good Set Query
提出コード
code:cpp
template<class Abel> struct UnionFind {
vector<int> par;
vector<int> rank;
vector<Abel> diff_weight;
UnionFind(int n = 1, Abel SUM_UNITY = 0) {
init(n, SUM_UNITY);
}
void init(int n = 1, Abel SUM_UNITY = 0) {
par.resize(n); rank.resize(n); diff_weight.resize(n);
for (int i = 0; i < n; ++i) pari = i, ranki = 0, diff_weighti = SUM_UNITY;
}
int root(int x) {
if (parx == x) {
return x;
}
else {
int r = root(parx);
diff_weightx += diff_weight[parx];
return parx = r;
}
}
Abel weight(int x) {
root(x);
return diff_weightx;
}
bool issame(int x, int y) {
return root(x) == root(y);
}
bool merge(int x, int y, Abel w) {
w += weight(x); w -= weight(y);
x = root(x); y = root(y);
if (x == y) return false;
if (rankx < ranky) swap(x, y), w = -w;
if (rankx == ranky) ++rankx;
pary = x;
diff_weighty = w;
return true;
}
Abel diff(int x, int y) {
return weight(y) - weight(x);
}
};
bool solve();
void _main(){
int testcase = 1;
// cin >> testcase;
for(;testcase--;){
if(solve()){
// O("Yes");
}
else{
// O("No");
}
}
}
bool solve(){
LL(n,q);
UnionFind<ll>uf(n);
vector<ll>ans;
rep(i,q){
LL(a,b,c);
a--,b--;
if(uf.issame(a,b)){
if(uf.diff(a,b)==c){
ans.pb(i+1);
}
}
else{
uf.merge(a,b,c);
ans.pb(i+1);
}
}
O(ans);
return true;
}