Test/AtCoder/abc340_g.test.cpp

• View this file on GitHub
• Last update: 2024-12-07 19:29:25+09:00
• Problem: https://atcoder.jp/contests/abc340/tasks/abc340_g

Depends on

• Src/Algebra/Monoid/ChminMonoid.hpp
• Sparse Table (Src/DataStructure/SparseTable/SparseTable.hpp)
• Auxiliary Tree (指定された頂点たちの最小共通祖先関係を保って木を圧縮してできる補助的な木) (Src/Graph/Tree/AuxiliaryTree.hpp)
• Lowest Common Ancestor (Src/Graph/Tree/LowestCommonAncestor.hpp)
• 標準データ型のエイリアス (Src/Template/TypeAlias.hpp)

Code

#define PROBLEM "https://atcoder.jp/contests/abc340/tasks/abc340_g"

#include "../../Src/Template/TypeAlias.hpp"
#include "../../Src/Graph/Tree/AuxiliaryTree.hpp"
#include "atcoder/modint"

#include <array>
#include <algorithm>
#include <iostream>
#include <set>

using namespace zawa;
using mint = atcoder::modint998244353;

int main() {
    int N;
    std::cin >> N;
    std::vector<std::vector<int>> A(N);
    std::vector<int> col(N);
    for (int i{} ; i < N ; i++) {
        std::cin >> col[i];
        col[i]--;
        A[col[i]].push_back(i);
    }
    std::vector<std::vector<int>> g(N);
    for (int i{} ; i < N - 1 ; i++) {
        int u, v;
        std::cin >> u >> v;
        u--; v--;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    AuxiliaryTree AT{g};
    mint ans{};
    for (int c{} ; c < N ; c++) {
        if (A[c].empty()) {
            continue;
        }
        int r{(int)AT.construct(A[c])};
        auto dfs{[&](auto dfs, int v, int p) -> mint {
            std::array<mint, 3> dp{ mint{1}, mint{0}, mint{0} };
            for (auto x : AT[v]) {
                if ((int)x == p) continue;
                std::array<mint, 3> next{ dp };
                mint chval{dfs(dfs, x, v)};
                for (int i{} ; i < 3 ; i++) {
                    next[std::min(2, i + 1)] += dp[i] * chval;
                }
                dp = next;
            }
            ans += dp[2];
            if (col[v] == c) {
                dp[1]++;
                ans += dp[1];
            }
            return dp[1] + dp[2];
        }};
        dfs(dfs, r, -1);
    }
    std::cout << ans.val() << '\n';
}

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.13.5/x64/lib/python3.13/site-packages/onlinejudge_verify/documentation/build.py", line 71, in _render_source_code_stat
    bundled_code = language.bundle(stat.path, basedir=basedir, options={'include_paths': [basedir]}).decode()
                   ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/hostedtoolcache/Python/3.13.5/x64/lib/python3.13/site-packages/onlinejudge_verify/languages/cplusplus.py", line 187, in bundle
    bundler.update(path)
    ~~~~~~~~~~~~~~^^^^^^
  File "/opt/hostedtoolcache/Python/3.13.5/x64/lib/python3.13/site-packages/onlinejudge_verify/languages/cplusplus_bundle.py", line 401, in update
    self.update(self._resolve(pathlib.Path(included), included_from=path))
                ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/hostedtoolcache/Python/3.13.5/x64/lib/python3.13/site-packages/onlinejudge_verify/languages/cplusplus_bundle.py", line 260, in _resolve
    raise BundleErrorAt(path, -1, "no such header")
onlinejudge_verify.languages.cplusplus_bundle.BundleErrorAt: atcoder/modint: line -1: no such header

Back to top page