Submission #47553082

---

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://scrapbox.io/sifi-algo/Quick_Find
struct QuickFind {
    std::vector<Num> leaders_;
    std::vector<std::vector<Num>> members_;
    std::vector<Num> potentials_;

    QuickFind(Num n_nodes) :
        leaders_(n_nodes, 0), members_(n_nodes), potentials_(n_nodes, 0) {
        std::iota(leaders_.begin(), leaders_.end(), Num{0});
        for(Num i{0}; i<n_nodes; ++i) {
            members_.at(i).push_back(i);
        }
    }

    Num leader(Num x) {
        return leaders_.at(x);
    }

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

        auto leader_x = leaders_.at(x);
        auto leader_y = leaders_.at(y);

        if (leader_x == leader_y) {
            return;
        }

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

        for(const auto& v : members_.at(leader_y)) {
            leaders_.at(v) = leader_x;
            potentials_.at(v) += p;
        }

        members_.at(leader_x).insert(
            members_.at(leader_x).end(),
            members_.at(leader_y).begin(), members_.at(leader_y).end());

        std::vector<Num> zero;
        std::swap(zero, members_.at(leader_y));
        return;
    }

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

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

    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;
    QuickFind 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