Submission #47550446

---

Source Code Expand

#include <bits/stdc++.h>
#include <atcoder/modint>

namespace {
using ModInt [[maybe_unused]] = atcoder::modint998244353;
using Num [[maybe_unused]] = long long int;
using Vec [[maybe_unused]] = std::vector<Num>;
using Set [[maybe_unused]] = std::set<Num>;
using Mset [[maybe_unused]] = std::multiset<Num>;
using Edges [[maybe_unused]] = std::vector<std::vector<Num>>;

template<typename T>
using Q [[maybe_unused]] = std::queue<T>;

template<typename T>
using PQ [[maybe_unused]] = std::priority_queue<T, std::vector<T>, std::greater<T>>;
}

// Based on
// https://qiita.com/drken/items/cce6fc5c579051e64fab
struct WeightedUnionFind {
    std::vector<Num> parents_;
	std::vector<Num> sizes_;
	std::vector<Num> potentials_;

	WeightedUnionFind(Num n_nodes) :
        parents_(n_nodes, 0), sizes_(n_nodes, 0), potentials_(n_nodes, 0) {
        std::iota(parents_.begin(), parents_.end(), Num{0});
	}

	Num leader(Num x) {
		if (parents_.at(x) == x) {
			return x;
		}

        const auto root = leader(parents_.at(x));
        potentials_.at(x) += potentials_.at(parents_.at(x));
        parents_.at(x) = root;
        return root;
	}

	Num potential(Num x) {
		leader(x);
		return potentials_.at(x);
	}

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

	void merge(Num x, Num y, Num potential) {
        auto p = potential;
		p += potentials_.at(x);
        p -= potentials_.at(y);

		auto leader_x = leader(x);
		auto leader_y = leader(y);
		if (leader_x == leader_y) {
            return;
        }

		if (sizes_.at(leader_x) < sizes_.at(leader_y)) {
            std::swap(leader_x, leader_y);
            p = -p;
        }

		if (sizes_.at(leader_x) == sizes_.at(leader_y)) {
            sizes_.at(leader_x) += 1;
        }

		parents_.at(leader_y) = leader_x;
		potentials_.at(leader_y) = p;
		return;
	}

	Num diff(Num x, Num y) {
		return potentials_.at(y) - potentials_.at(x);
	}
};

void solve(std::istream& is, std::ostream& os) {
    Num n, q;
    is >> n >> q;

    Vec answer;
    WeightedUnionFind dsu(n);

    for(Num i{1}; i<=q; ++i) {
        Num a, b, d;
        is >> b >> a >> d;
        --a;
        --b;

        bool found = dsu.same(a, b);
        if (!found || (dsu.diff(a, b) == d)) {
            answer.push_back(i);
        }

        if (!found) {
            dsu.merge(a, b, d);
        }
    }

    const auto size = answer.size();
    for(size_t i{0}; i<size; ++i) {
        os << answer.at(i);
        if ((i+1) == size) {
            os << "\n";
        } else {
            os << " ";
        }
    }
}

int main(void) {
    solve(std::cin, std::cout);
    return 0;
}

Submission Info

Judge Result