ABC328 F問題（Good Set Query）を解く

#include <bits/stdc++.h>
using namespace std;

typedef long long int ll;

class Weighted_UnionFind {
private:
	int n;
	std::vector<int> par;
	std::vector<ll> diff_weight;

public:
	Weighted_UnionFind(int _n) {
		n = _n;
		par.resize(n, -1);
		diff_weight.resize(n, 0);
	}

	int leader(int x) {
		if (par[x] < 0) {
			return x;
		} else {
			int r = leader(par[x]);
			diff_weight[x] += diff_weight[par[x]];
			return par[x] = r;
		}
	}

	ll weight(int x) {
		leader(x);
		return diff_weight[x];
	}

	ll diff(int x, int y) {
		return weight(y) - weight(x);
	}

	void merge(int x, int y, ll w) {
		w += weight(x);
		w -= weight(y);
		x = leader(x);
		y = leader(y);
		if (x != y) {
			if (-par[x] < -par[y]) {
				std::swap(x, y);
				w = -w;
			}
			par[x] += par[y];
			par[y] = x;
			diff_weight[y] = w;
		}
	}

	bool same(int x, int y) {
		return leader(x) == leader(y);
	}

	int size(int x) {
		return -par[leader(x)];
	}

	std::vector<std::vector<int>> groups() {
		std::vector<std::vector<int>> member(n);
		for (int i = 0; i < n; ++i) {
			member[leader(i)].push_back(i);
		}

		std::vector<std::vector<int>> result;
		for (int i = 0; i < n; ++i) {
			if (!member[i].empty()) {
				result.push_back(member[i]);
			}
		}

		return result;
	}
};

int main()
{
	int n, q;
	cin >> n >> q;
	Weighted_UnionFind uf(n + 1);
	vector<int> result;
	for (int i = 0; i < q; ++i) {
		int a, b;
		ll d;
		cin >> a >> b >> d;
		if (uf.same(a, b)) {
			if (uf.diff(b, a) != d) {
				continue;
			}
		} else {
			uf.merge(b, a, d);
		}
		result.push_back(i + 1);
	}

	for (int i = 0; i < result.size(); ++i) {
		cout << result[i] << " \n"[i == result.size() - 1];
	}

	return 0;
}