Submission #47551491

---

Source Code Expand

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

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

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

    Edges edges(n);
    std::map<std::pair<Num, Num>, Num> weights;
    std::vector<std::tuple<Num, Num, Num>> queries(q);
    atcoder::dsu tree(n);

    for(Num i{0}; i<q; ++i) {
        Num a, b, d;
        is >> b >> a >> d;
        --a;
        --b;
        queries.at(i) = std::make_tuple(a, b, d);

        if (!tree.same(a, b)) {
            tree.merge(a, b);
            edges[a].push_back(b);
            edges[b].push_back(a);
            weights[std::make_pair(a,b)] = d;
            weights[std::make_pair(b,a)] = -d;
        }
    }

    std::vector<std::optional<Num>> distances(n);
    for(Num i{0}; i<n; ++i) {
        if (distances.at(i).has_value()) {
            continue;
        }

        Q<Num> q;
        auto current = tree.leader(i);
        distances.at(current) = 0;
        q.push(current);

        while(!q.empty()) {
            const auto current = q.front();
            q.pop();

            for(const auto& next : edges[current]) {
                if (distances.at(next).has_value()) {
                    continue;
                }
                distances.at(next) = distances.at(current).value() + weights.at(std::make_pair(current, next));
                q.push(next);
            }
        }
    }

    Vec answer;
    for(Num i{0}; i<q; ++i) {
        const auto& [a, b, d] = queries.at(i);
        const auto dist = distances.at(b).value() - distances.at(a).value();
        if (dist == d) {
            answer.push_back(i+1);
        }
    }

    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