AtCoder Beginner Contest 328 F - Good Set Query

F - Good Set Query

由题意：我们其实不难看出这题就是要维护Xai，Xbi之间的差值关系（Xai - Xbi = di可以理解为ai->bi的距离为di， sz[bi] = sz[ai] + di），自然而然的想到加权并查集。pre维护根结点，sz维护与根结点的距离。

#include <bits/stdc++.h>
#define SQR(x) fixed << setprecision(x)
using namespace std;

using ll = long long;

const int N = 2e5 + 10;

ll pre[N], sz[N];

int find(int f) {
	if(pre[f] == f) return f;
	int fi = find(pre[f]);
	sz[f] += sz[pre[f]];  //根结点更新，sz值也要更新
	return pre[f] = fi;
}

bool join(int a, int b, ll d) {
	int fa = find(a), fb = find(b);
	if(fa == fb) return sz[b] - sz[a] == d; //如果他们在同一个块里，则只需要判断
	//之间的差值是否为d
	//否则合并
	pre[fa] = fb;
	sz[fa] = sz[b] - sz[a] - d;  //将一个父结点的sz值改变，这里是fa
	//（因为它的父节点被修改），可以理解为fa->a->b->fb ： -sz[a]-d+sz[b]
	return true;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
	
	int n, q;
	cin >> n >> q;

	for (int i = 1; i <= n; i++) pre[i] = i;
	for (int i = 1; i <= q; i++) {
		int a, b, d;
		cin >> a >> b >> d;
		if(join(a, b, d)) {
			cout << i << ' ';
		}
	}    
	cout << "\n";

    return 0;
}