This documentation is automatically generated by online-judge-tools/verification-helper
View the Project on GitHub hitonanode/cplib-cpp
#include "../../convolution/ntt.hpp" #include "../../modint.hpp" #include "../../number/bare_mod_algebra.hpp" #include "../frequency_table_of_tree_distance.hpp" #include <iostream> #include <vector> #define PROBLEM "https://judge.yosupo.jp/problem/frequency_table_of_tree_distance" using namespace std; int main() { cin.tie(nullptr), ios::sync_with_stdio(false); int N; cin >> N; vector<vector<int>> to(N); for (int i = 0; i < N - 1; i++) { int s, t; cin >> s >> t; to[s].emplace_back(t), to[t].emplace_back(s); } frequency_table_of_tree_distance solver(to); using mint1 = ModInt<998244353>; using mint2 = ModInt<167772161>; const vector<mint1> ret1 = frequency_table_of_tree_distance(to).solve<mint1, nttconv<mint1>>(std::vector<mint1>(N, 1)); const vector<mint2> ret2 = frequency_table_of_tree_distance(to).solve<mint2, nttconv<mint2>>(std::vector<mint2>(N, 1)); for (int i = 1; i < N; i++) { auto v = crt<long long>({ret1[i].val(), ret2[i].val()}, {mint1::mod(), mint2::mod()}); cout << v.first << ' '; } cout << '\n'; }
#line 2 "modint.hpp" #include <cassert> #include <iostream> #include <set> #include <vector> template <int md> struct ModInt { using lint = long long; constexpr static int mod() { return md; } static int get_primitive_root() { static int primitive_root = 0; if (!primitive_root) { primitive_root = [&]() { std::set<int> fac; int v = md - 1; for (lint i = 2; i * i <= v; i++) while (v % i == 0) fac.insert(i), v /= i; if (v > 1) fac.insert(v); for (int g = 1; g < md; g++) { bool ok = true; for (auto i : fac) if (ModInt(g).pow((md - 1) / i) == 1) { ok = false; break; } if (ok) return g; } return -1; }(); } return primitive_root; } int val_; int val() const noexcept { return val_; } constexpr ModInt() : val_(0) {} constexpr ModInt &_setval(lint v) { return val_ = (v >= md ? v - md : v), *this; } constexpr ModInt(lint v) { _setval(v % md + md); } constexpr explicit operator bool() const { return val_ != 0; } constexpr ModInt operator+(const ModInt &x) const { return ModInt()._setval((lint)val_ + x.val_); } constexpr ModInt operator-(const ModInt &x) const { return ModInt()._setval((lint)val_ - x.val_ + md); } constexpr ModInt operator*(const ModInt &x) const { return ModInt()._setval((lint)val_ * x.val_ % md); } constexpr ModInt operator/(const ModInt &x) const { return ModInt()._setval((lint)val_ * x.inv().val() % md); } constexpr ModInt operator-() const { return ModInt()._setval(md - val_); } constexpr ModInt &operator+=(const ModInt &x) { return *this = *this + x; } constexpr ModInt &operator-=(const ModInt &x) { return *this = *this - x; } constexpr ModInt &operator*=(const ModInt &x) { return *this = *this * x; } constexpr ModInt &operator/=(const ModInt &x) { return *this = *this / x; } friend constexpr ModInt operator+(lint a, const ModInt &x) { return ModInt(a) + x; } friend constexpr ModInt operator-(lint a, const ModInt &x) { return ModInt(a) - x; } friend constexpr ModInt operator*(lint a, const ModInt &x) { return ModInt(a) * x; } friend constexpr ModInt operator/(lint a, const ModInt &x) { return ModInt(a) / x; } constexpr bool operator==(const ModInt &x) const { return val_ == x.val_; } constexpr bool operator!=(const ModInt &x) const { return val_ != x.val_; } constexpr bool operator<(const ModInt &x) const { return val_ < x.val_; } // To use std::map<ModInt, T> friend std::istream &operator>>(std::istream &is, ModInt &x) { lint t; return is >> t, x = ModInt(t), is; } constexpr friend std::ostream &operator<<(std::ostream &os, const ModInt &x) { return os << x.val_; } constexpr ModInt pow(lint n) const { ModInt ans = 1, tmp = *this; while (n) { if (n & 1) ans *= tmp; tmp *= tmp, n >>= 1; } return ans; } static constexpr int cache_limit = std::min(md, 1 << 21); static std::vector<ModInt> facs, facinvs, invs; constexpr static void _precalculation(int N) { const int l0 = facs.size(); if (N > md) N = md; if (N <= l0) return; facs.resize(N), facinvs.resize(N), invs.resize(N); for (int i = l0; i < N; i++) facs[i] = facs[i - 1] * i; facinvs[N - 1] = facs.back().pow(md - 2); for (int i = N - 2; i >= l0; i--) facinvs[i] = facinvs[i + 1] * (i + 1); for (int i = N - 1; i >= l0; i--) invs[i] = facinvs[i] * facs[i - 1]; } constexpr ModInt inv() const { if (this->val_ < cache_limit) { if (facs.empty()) facs = {1}, facinvs = {1}, invs = {0}; while (this->val_ >= int(facs.size())) _precalculation(facs.size() * 2); return invs[this->val_]; } else { return this->pow(md - 2); } } constexpr ModInt fac() const { while (this->val_ >= int(facs.size())) _precalculation(facs.size() * 2); return facs[this->val_]; } constexpr ModInt facinv() const { while (this->val_ >= int(facs.size())) _precalculation(facs.size() * 2); return facinvs[this->val_]; } constexpr ModInt doublefac() const { lint k = (this->val_ + 1) / 2; return (this->val_ & 1) ? ModInt(k * 2).fac() / (ModInt(2).pow(k) * ModInt(k).fac()) : ModInt(k).fac() * ModInt(2).pow(k); } constexpr ModInt nCr(int r) const { if (r < 0 or this->val_ < r) return ModInt(0); return this->fac() * (*this - r).facinv() * ModInt(r).facinv(); } constexpr ModInt nPr(int r) const { if (r < 0 or this->val_ < r) return ModInt(0); return this->fac() * (*this - r).facinv(); } static ModInt binom(int n, int r) { static long long bruteforce_times = 0; if (r < 0 or n < r) return ModInt(0); if (n <= bruteforce_times or n < (int)facs.size()) return ModInt(n).nCr(r); r = std::min(r, n - r); ModInt ret = ModInt(r).facinv(); for (int i = 0; i < r; ++i) ret *= n - i; bruteforce_times += r; return ret; } // Multinomial coefficient, (k_1 + k_2 + ... + k_m)! / (k_1! k_2! ... k_m!) // Complexity: O(sum(ks)) template <class Vec> static ModInt multinomial(const Vec &ks) { ModInt ret{1}; int sum = 0; for (int k : ks) { assert(k >= 0); ret *= ModInt(k).facinv(), sum += k; } return ret * ModInt(sum).fac(); } // Catalan number, C_n = binom(2n, n) / (n + 1) // C_0 = 1, C_1 = 1, C_2 = 2, C_3 = 5, C_4 = 14, ... // https://oeis.org/A000108 // Complexity: O(n) static ModInt catalan(int n) { if (n < 0) return ModInt(0); return ModInt(n * 2).fac() * ModInt(n + 1).facinv() * ModInt(n).facinv(); } ModInt sqrt() const { if (val_ == 0) return 0; if (md == 2) return val_; if (pow((md - 1) / 2) != 1) return 0; ModInt b = 1; while (b.pow((md - 1) / 2) == 1) b += 1; int e = 0, m = md - 1; while (m % 2 == 0) m >>= 1, e++; ModInt x = pow((m - 1) / 2), y = (*this) * x * x; x *= (*this); ModInt z = b.pow(m); while (y != 1) { int j = 0; ModInt t = y; while (t != 1) j++, t *= t; z = z.pow(1LL << (e - j - 1)); x *= z, z *= z, y *= z; e = j; } return ModInt(std::min(x.val_, md - x.val_)); } }; template <int md> std::vector<ModInt<md>> ModInt<md>::facs = {1}; template <int md> std::vector<ModInt<md>> ModInt<md>::facinvs = {1}; template <int md> std::vector<ModInt<md>> ModInt<md>::invs = {0}; using ModInt998244353 = ModInt<998244353>; // using mint = ModInt<998244353>; // using mint = ModInt<1000000007>; #line 3 "convolution/ntt.hpp" #include <algorithm> #include <array> #line 7 "convolution/ntt.hpp" #include <tuple> #line 9 "convolution/ntt.hpp" // CUT begin // Integer convolution for arbitrary mod // with NTT (and Garner's algorithm) for ModInt / ModIntRuntime class. // We skip Garner's algorithm if `skip_garner` is true or mod is in `nttprimes`. // input: a (size: n), b (size: m) // return: vector (size: n + m - 1) template <typename MODINT> std::vector<MODINT> nttconv(std::vector<MODINT> a, std::vector<MODINT> b, bool skip_garner); constexpr int nttprimes[3] = {998244353, 167772161, 469762049}; // Integer FFT (Fast Fourier Transform) for ModInt class // (Also known as Number Theoretic Transform, NTT) // is_inverse: inverse transform // ** Input size must be 2^n ** template <typename MODINT> void ntt(std::vector<MODINT> &a, bool is_inverse = false) { int n = a.size(); if (n == 1) return; static const int mod = MODINT::mod(); static const MODINT root = MODINT::get_primitive_root(); assert(__builtin_popcount(n) == 1 and (mod - 1) % n == 0); static std::vector<MODINT> w{1}, iw{1}; for (int m = w.size(); m < n / 2; m *= 2) { MODINT dw = root.pow((mod - 1) / (4 * m)), dwinv = 1 / dw; w.resize(m * 2), iw.resize(m * 2); for (int i = 0; i < m; i++) w[m + i] = w[i] * dw, iw[m + i] = iw[i] * dwinv; } if (!is_inverse) { for (int m = n; m >>= 1;) { for (int s = 0, k = 0; s < n; s += 2 * m, k++) { for (int i = s; i < s + m; i++) { MODINT x = a[i], y = a[i + m] * w[k]; a[i] = x + y, a[i + m] = x - y; } } } } else { for (int m = 1; m < n; m *= 2) { for (int s = 0, k = 0; s < n; s += 2 * m, k++) { for (int i = s; i < s + m; i++) { MODINT x = a[i], y = a[i + m]; a[i] = x + y, a[i + m] = (x - y) * iw[k]; } } } int n_inv = MODINT(n).inv().val(); for (auto &v : a) v *= n_inv; } } template <int MOD> std::vector<ModInt<MOD>> nttconv_(const std::vector<int> &a, const std::vector<int> &b) { int sz = a.size(); assert(a.size() == b.size() and __builtin_popcount(sz) == 1); std::vector<ModInt<MOD>> ap(sz), bp(sz); for (int i = 0; i < sz; i++) ap[i] = a[i], bp[i] = b[i]; ntt(ap, false); if (a == b) bp = ap; else ntt(bp, false); for (int i = 0; i < sz; i++) ap[i] *= bp[i]; ntt(ap, true); return ap; } long long garner_ntt_(int r0, int r1, int r2, int mod) { using mint2 = ModInt<nttprimes[2]>; static const long long m01 = 1LL * nttprimes[0] * nttprimes[1]; static const long long m0_inv_m1 = ModInt<nttprimes[1]>(nttprimes[0]).inv().val(); static const long long m01_inv_m2 = mint2(m01).inv().val(); int v1 = (m0_inv_m1 * (r1 + nttprimes[1] - r0)) % nttprimes[1]; auto v2 = (mint2(r2) - r0 - mint2(nttprimes[0]) * v1) * m01_inv_m2; return (r0 + 1LL * nttprimes[0] * v1 + m01 % mod * v2.val()) % mod; } template <typename MODINT> std::vector<MODINT> nttconv(std::vector<MODINT> a, std::vector<MODINT> b, bool skip_garner) { if (a.empty() or b.empty()) return {}; int sz = 1, n = a.size(), m = b.size(); while (sz < n + m) sz <<= 1; if (sz <= 16) { std::vector<MODINT> ret(n + m - 1); for (int i = 0; i < n; i++) { for (int j = 0; j < m; j++) ret[i + j] += a[i] * b[j]; } return ret; } int mod = MODINT::mod(); if (skip_garner or std::find(std::begin(nttprimes), std::end(nttprimes), mod) != std::end(nttprimes)) { a.resize(sz), b.resize(sz); if (a == b) { ntt(a, false); b = a; } else { ntt(a, false), ntt(b, false); } for (int i = 0; i < sz; i++) a[i] *= b[i]; ntt(a, true); a.resize(n + m - 1); } else { std::vector<int> ai(sz), bi(sz); for (int i = 0; i < n; i++) ai[i] = a[i].val(); for (int i = 0; i < m; i++) bi[i] = b[i].val(); auto ntt0 = nttconv_<nttprimes[0]>(ai, bi); auto ntt1 = nttconv_<nttprimes[1]>(ai, bi); auto ntt2 = nttconv_<nttprimes[2]>(ai, bi); a.resize(n + m - 1); for (int i = 0; i < n + m - 1; i++) a[i] = garner_ntt_(ntt0[i].val(), ntt1[i].val(), ntt2[i].val(), mod); } return a; } template <typename MODINT> std::vector<MODINT> nttconv(const std::vector<MODINT> &a, const std::vector<MODINT> &b) { return nttconv<MODINT>(a, b, false); } #line 5 "number/bare_mod_algebra.hpp" #include <utility> #line 7 "number/bare_mod_algebra.hpp" // CUT begin // Solve ax+by=gcd(a, b) template <class Int> Int extgcd(Int a, Int b, Int &x, Int &y) { Int d = a; if (b != 0) { d = extgcd(b, a % b, y, x), y -= (a / b) * x; } else { x = 1, y = 0; } return d; } // Calculate a^(-1) (MOD m) s if gcd(a, m) == 1 // Calculate x s.t. ax == gcd(a, m) MOD m template <class Int> Int mod_inverse(Int a, Int m) { Int x, y; extgcd<Int>(a, m, x, y); x %= m; return x + (x < 0) * m; } // Require: 1 <= b // return: (g, x) s.t. g = gcd(a, b), xa = g MOD b, 0 <= x < b/g template <class Int> /* constexpr */ std::pair<Int, Int> inv_gcd(Int a, Int b) { a %= b; if (a < 0) a += b; if (a == 0) return {b, 0}; Int s = b, t = a, m0 = 0, m1 = 1; while (t) { Int u = s / t; s -= t * u, m0 -= m1 * u; auto tmp = s; s = t, t = tmp, tmp = m0, m0 = m1, m1 = tmp; } if (m0 < 0) m0 += b / s; return {s, m0}; } template <class Int> /* constexpr */ std::pair<Int, Int> crt(const std::vector<Int> &r, const std::vector<Int> &m) { assert(r.size() == m.size()); int n = int(r.size()); // Contracts: 0 <= r0 < m0 Int r0 = 0, m0 = 1; for (int i = 0; i < n; i++) { assert(1 <= m[i]); Int r1 = r[i] % m[i], m1 = m[i]; if (r1 < 0) r1 += m1; if (m0 < m1) { std::swap(r0, r1); std::swap(m0, m1); } if (m0 % m1 == 0) { if (r0 % m1 != r1) return {0, 0}; continue; } Int g, im; std::tie(g, im) = inv_gcd<Int>(m0, m1); Int u1 = m1 / g; if ((r1 - r0) % g) return {0, 0}; Int x = (r1 - r0) / g % u1 * im % u1; r0 += x * m0; m0 *= u1; if (r0 < 0) r0 += m0; } return {r0, m0}; } // 蟻本 P.262 // 中国剰余定理を利用して,色々な素数で割った余りから元の値を復元 // 連立線形合同式 A * x = B mod M の解 // Requirement: M[i] > 0 // Output: x = first MOD second (if solution exists), (0, 0) (otherwise) template <class Int> std::pair<Int, Int> linear_congruence(const std::vector<Int> &A, const std::vector<Int> &B, const std::vector<Int> &M) { Int r = 0, m = 1; assert(A.size() == M.size()); assert(B.size() == M.size()); for (int i = 0; i < (int)A.size(); i++) { assert(M[i] > 0); const Int ai = A[i] % M[i]; Int a = ai * m, b = B[i] - ai * r, d = std::__gcd(M[i], a); if (b % d != 0) { return std::make_pair(0, 0); // 解なし } Int t = b / d * mod_inverse<Int>(a / d, M[i] / d) % (M[i] / d); r += m * t; m *= M[i] / d; } return std::make_pair((r < 0 ? r + m : r), m); } template <class Int = int, class Long = long long> Int pow_mod(Int x, long long n, Int md) { static_assert(sizeof(Int) * 2 <= sizeof(Long), "Watch out for overflow"); if (md == 1) return 0; Int ans = 1; while (n > 0) { if (n & 1) ans = (Long)ans * x % md; x = (Long)x * x % md; n >>= 1; } return ans; } #line 5 "tree/centroid_decomposition.hpp" // CUT begin /* (Recursive) Centroid Decomposition Verification: Codeforces #190 Div.1 C https://codeforces.com/contest/321/submission/59093583 fix_root(int r): Build information of the tree which `r` belongs to. detect_centroid(int r): Enumerate centroid(s) of the tree which `r` belongs to. */ struct CentroidDecomposition { int NO_PARENT = -1; int V; int E; std::vector<std::vector<std::pair<int, int>>> to; // (node_id, edge_id) std::vector<int> par; // parent node_id par[root] = -1 std::vector<std::vector<int>> chi; // children id's std::vector<int> subtree_size; // size of each subtree std::vector<int> available_edge; // If 0, ignore the corresponding edge. CentroidDecomposition(int v = 0) : V(v), E(0), to(v), par(v, NO_PARENT), chi(v), subtree_size(v) {} CentroidDecomposition(const std::vector<std::vector<int>> &to_) : CentroidDecomposition(to_.size()) { for (int i = 0; i < V; i++) { for (auto j : to_[i]) { if (i < j) { add_edge(i, j); } } } } void add_edge(int v1, int v2) { to[v1].emplace_back(v2, E), to[v2].emplace_back(v1, E), E++; available_edge.emplace_back(1); } int _dfs_fixroot(int now, int prv) { chi[now].clear(), subtree_size[now] = 1; for (auto nxt : to[now]) { if (nxt.first != prv and available_edge[nxt.second]) { par[nxt.first] = now, chi[now].push_back(nxt.first); subtree_size[now] += _dfs_fixroot(nxt.first, now); } } return subtree_size[now]; } void fix_root(int root) { par[root] = NO_PARENT; _dfs_fixroot(root, -1); } //// Centroid Decpmposition //// std::vector<int> centroid_cand_tmp; void _dfs_detect_centroids(int now, int prv, int n) { bool is_centroid = true; for (auto nxt : to[now]) { if (nxt.first != prv and available_edge[nxt.second]) { _dfs_detect_centroids(nxt.first, now, n); if (subtree_size[nxt.first] > n / 2) is_centroid = false; } } if (n - subtree_size[now] > n / 2) is_centroid = false; if (is_centroid) centroid_cand_tmp.push_back(now); } std::pair<int, int> detect_centroids(int r) { // ([centroid_node_id1], ([centroid_node_id2]|-1)) centroid_cand_tmp.clear(); while (par[r] != NO_PARENT) r = par[r]; int n = subtree_size[r]; _dfs_detect_centroids(r, -1, n); if (centroid_cand_tmp.size() == 1) return std::make_pair(centroid_cand_tmp[0], -1); else return std::make_pair(centroid_cand_tmp[0], centroid_cand_tmp[1]); } std::vector<int> _cd_vertices; void _centroid_decomposition(int now) { fix_root(now); now = detect_centroids(now).first; _cd_vertices.emplace_back(now); /* do something */ for (auto p : to[now]) { int nxt, eid; std::tie(nxt, eid) = p; if (available_edge[eid] == 0) continue; available_edge[eid] = 0; _centroid_decomposition(nxt); } } std::vector<int> centroid_decomposition(int x) { _cd_vertices.clear(); _centroid_decomposition(x); return _cd_vertices; } }; #line 5 "tree/frequency_table_of_tree_distance.hpp" struct frequency_table_of_tree_distance { std::vector<std::vector<int>> tos; std::vector<int> cd; std::vector<std::pair<int, int>> tmp; std::vector<int> alive; void _dfs(int now, int prv, int depth) { // if (int(tmp.size()) <= depth) tmp.resize(depth + 1, 0); // tmp[depth]++; tmp.emplace_back(now, depth); for (auto nxt : tos[now]) { if (alive[nxt] and nxt != prv) _dfs(nxt, now, depth + 1); } } std::vector<std::pair<int, int>> cnt_dfs(int root) { return tmp.clear(), _dfs(root, -1, 0), tmp; } frequency_table_of_tree_distance(const std::vector<std::vector<int>> &to) { tos = to; cd = CentroidDecomposition(to).centroid_decomposition(0); } template <class S, std::vector<S> (*conv)(const std::vector<S> &, const std::vector<S> &)> std::vector<S> solve(const std::vector<S> &weight) { alive.assign(tos.size(), 1); std::vector<S> ret(tos.size()); std::vector<S> v; for (auto root : cd) { std::vector<std::vector<S>> vv; alive[root] = 0; for (auto nxt : tos[root]) { if (!alive[nxt]) continue; v.clear(); for (auto p : cnt_dfs(nxt)) { while (int(v.size()) <= p.second) v.push_back(S(0)); v[p.second] += weight[p.first]; } for (int i = 0; i < int(v.size()); i++) ret[i + 1] += v[i] * weight[root]; vv.emplace_back(v); } std::sort(vv.begin(), vv.end(), [&](const std::vector<S> &l, const std::vector<S> &r) { return l.size() < r.size(); }); for (size_t j = 1; j < vv.size(); j++) { const std::vector<S> c = conv(vv[j - 1], vv[j]); for (size_t i = 0; i < c.size(); i++) ret[i + 2] += c[i]; for (size_t i = 0; i < vv[j - 1].size(); i++) vv[j][i] += vv[j - 1][i]; } tos[root].clear(); } return ret; } }; #line 7 "tree/test/frequency_table_of_tree_distance_ntt.test.cpp" #define PROBLEM "https://judge.yosupo.jp/problem/frequency_table_of_tree_distance" using namespace std; int main() { cin.tie(nullptr), ios::sync_with_stdio(false); int N; cin >> N; vector<vector<int>> to(N); for (int i = 0; i < N - 1; i++) { int s, t; cin >> s >> t; to[s].emplace_back(t), to[t].emplace_back(s); } frequency_table_of_tree_distance solver(to); using mint1 = ModInt<998244353>; using mint2 = ModInt<167772161>; const vector<mint1> ret1 = frequency_table_of_tree_distance(to).solve<mint1, nttconv<mint1>>(std::vector<mint1>(N, 1)); const vector<mint2> ret2 = frequency_table_of_tree_distance(to).solve<mint2, nttconv<mint2>>(std::vector<mint2>(N, 1)); for (int i = 1; i < N; i++) { auto v = crt<long long>({ret1[i].val(), ret2[i].val()}, {mint1::mod(), mint2::mod()}); cout << v.first << ' '; } cout << '\n'; }