test/src/graph/enumerate_triangles/enumerate_triangles.test.cpp

• View this file on GitHub
• Last update: 2022-07-05 15:09:50+09:00
• Problem: https://judge.yosupo.jp/problem/enumerate_triangles

Depends on

• 三角形列挙 (library/graph/enumerate_triangles.hpp)

Code

#define PROBLEM "https://judge.yosupo.jp/problem/enumerate_triangles"

#include <iostream>
#include <atcoder/modint>

using mint = atcoder::modint998244353;

std::istream& operator>>(std::istream& in, mint &a) {
    long long e; in >> e; a = e;
    return in;
}

std::ostream& operator<<(std::ostream& out, const mint &a) {
    out << a.val();
    return out;
}

#include "library/graph/enumerate_triangles.hpp"

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int n, m;
    std::cin >> n >> m;

    std::vector<mint> x(n);
    for (auto &e : x) std::cin >> e;

    std::vector<std::vector<int>> g(n);
    for (int i = 0; i < m; ++i) {
        int u, v;
        std::cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    mint ans = 0;
    suisen::enumerate_triangles(g, [&](int i, int j, int k) { ans += x[i] * x[j] * x[k]; });
    std::cout << ans << std::endl;

    return 0;
}

#line 1 "test/src/graph/enumerate_triangles/enumerate_triangles.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/enumerate_triangles"

#include <iostream>
#include <atcoder/modint>

using mint = atcoder::modint998244353;

std::istream& operator>>(std::istream& in, mint &a) {
    long long e; in >> e; a = e;
    return in;
}

std::ostream& operator<<(std::ostream& out, const mint &a) {
    out << a.val();
    return out;
}

#line 1 "library/graph/enumerate_triangles.hpp"



#include <algorithm>
#include <vector>

namespace suisen {
    /**
     * Type Parameters
     * - TriangleConsumer : type of consumer function (int, int, int) -> void
     * 
     * Parameters
     * - std::vector<std::vector<int>> g : simple undirected graph
     * - TriangleConsumer fun
     * 
     * Requirements
     * - v in g[u] <=> u in g[v]
     * 
     * Complexity
     * - time : O(N + M^{3/2})
     * - space : O(N + M)
     */
    template <typename TriangleConsumer>
    void enumerate_triangles(std::vector<std::vector<int>> g, TriangleConsumer &&fun) {
        const int n = g.size();
        std::vector<int> cnt(n + 1);
        for (int i = 0; i < n; ++i) ++cnt[g[i].size() + 1];
        for (int i = 0; i < n; ++i) cnt[i + 1] += cnt[i];
        std::vector<int> ord(n), idx(n);
        for (int i = 0; i < n; ++i) ord[idx[i] = cnt[g[i].size()]++] = i;
        for (int i = 0; i < n; ++i) g[i].erase(std::remove_if(g[i].begin(), g[i].end(), [&](int j) { return idx[j] < idx[i]; }), g[i].end());
        std::vector<int8_t> exists(n, false);
        for (int i : ord) {
            for (int j : g[i]) exists[j] = true;
            for (int j : g[i]) for (int k : g[j]) if (exists[k]) fun(i, j, k);
            for (int j : g[i]) exists[j] = false;
        }
    }
} // namespace suisen


#line 19 "test/src/graph/enumerate_triangles/enumerate_triangles.test.cpp"

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int n, m;
    std::cin >> n >> m;

    std::vector<mint> x(n);
    for (auto &e : x) std::cin >> e;

    std::vector<std::vector<int>> g(n);
    for (int i = 0; i < m; ++i) {
        int u, v;
        std::cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }

    mint ans = 0;
    suisen::enumerate_triangles(g, [&](int i, int j, int k) { ans += x[i] * x[j] * x[k]; });
    std::cout << ans << std::endl;

    return 0;
}

Back to top page