sotanishy's code snippets for competitive programming
View the Project on GitHub sotanishy/cp-library-cpp
#define PROBLEM "https://judge.yosupo.jp/problem/frequency_table_of_tree_distance" #include "../../tree/centroid_decomposition.hpp" #include "../../convolution/fft.hpp" #include <bits/stdc++.h> using namespace std; using ll = long long; #define rep(i, s, t) for (int i = (int)(s); i < (int)(t); ++i) #define revrep(i, t, s) for (int i = (int)(t)-1; i >= (int)(s); --i) #define all(x) begin(x), end(x) template <typename T> bool chmax(T& a, const T& b) { return a < b ? (a = b, 1) : 0; } template <typename T> bool chmin(T& a, const T& b) { return a > b ? (a = b, 1) : 0; } int main() { ios_base::sync_with_stdio(false); cin.tie(nullptr); cout << fixed << setprecision(15); int N; cin >> N; vector<vector<int>> G(N); rep(i,0,N-1) { int a, b; cin >> a >> b; G[a].push_back(b); G[b].push_back(a); } vector<int> level, sz, par; tie(level, sz, par) = centroid_decomposition(G); vector<ll> ans(N); vector<double> cnt; auto dfs = [&](auto& dfs, int v, int p, int d, int l) -> void { if (cnt.size() <= d) cnt.push_back(0); ++cnt[d]; for (int c : G[v]) { if (c != p && level[c] > l) { dfs(dfs, c, v, d+1, l); } } }; auto calc = [&](int s) { vector<double> res(sz[s]); cnt.assign(1, 0); for (int c : G[s]) { if (level[c] > level[s]) { dfs(dfs, c, s, 1, level[s]); } } auto conv = convolution(cnt, cnt); rep(d,0,cnt.size()) { res[d] += cnt[d]; } rep(d,0,min(conv.size(), res.size())) { res[d] += conv[d]/2; } for (int c : G[s]) { if (level[c] > level[s]) { cnt.assign(1, 0); dfs(dfs, c, s, 1, level[s]); auto conv = convolution(cnt, cnt); rep(d,0,min(conv.size(), res.size())) { res[d] -= conv[d]/2; } } } rep(d,0,sz[s]) { ans[d] += (ll) round(res[d]); } }; rep(i,0,N) calc(i); rep(d,1,N) cout << ans[d] << (d < N-1 ? " " : "\n"); }
#line 1 "test/yosupo/frequency_table_of_tree_distance.test.cpp" #define PROBLEM "https://judge.yosupo.jp/problem/frequency_table_of_tree_distance" #line 2 "tree/centroid_decomposition.hpp" #include <tuple> #include <vector> std::tuple<std::vector<int>, std::vector<int>, std::vector<int>> centroid_decomposition(const std::vector<std::vector<int>>& G) { const int N = G.size(); std::vector<int> sz(N), level(N, -1), sz_comp(N), par(N); auto dfs_size = [&](auto& dfs_size, int v, int p) -> int { sz[v] = 1; for (int c : G[v]) { if (c != p && level[c] == -1) sz[v] += dfs_size(dfs_size, c, v); } return sz[v]; }; auto dfs_centroid = [&](auto& dfs_centroid, int v, int p, int n) -> int { for (int c : G[v]) { if (c != p && level[c] == -1 && sz[c] > n / 2) return dfs_centroid(dfs_centroid, c, v, n); } return v; }; auto decompose = [&](auto& decompose, int v, int k, int p) -> void { int n = dfs_size(dfs_size, v, -1); int s = dfs_centroid(dfs_centroid, v, -1, n); level[s] = k; sz_comp[s] = n; par[s] = p; for (int c : G[s]) { if (level[c] == -1) decompose(decompose, c, k + 1, s); } }; decompose(decompose, 0, 0, -1); return {level, sz_comp, par}; } #line 2 "convolution/fft.hpp" #include <bit> #include <complex> #include <numbers> #line 6 "convolution/fft.hpp" void fft(std::vector<std::complex<double>>& a) { const int n = a.size(); for (int m = n; m > 1; m >>= 1) { double ang = 2.0 * std::numbers::pi / m; std::complex<double> omega(cos(ang), sin(ang)); for (int s = 0; s < n / m; ++s) { std::complex<double> w(1, 0); for (int i = 0; i < m / 2; ++i) { auto l = a[s * m + i]; auto r = a[s * m + i + m / 2]; a[s * m + i] = l + r; a[s * m + i + m / 2] = (l - r) * w; w *= omega; } } } } void ifft(std::vector<std::complex<double>>& a) { const int n = a.size(); for (int m = 2; m <= n; m <<= 1) { double ang = -2.0 * std::numbers::pi / m; std::complex<double> omega(cos(ang), sin(ang)); for (int s = 0; s < n / m; ++s) { std::complex<double> w(1, 0); for (int i = 0; i < m / 2; ++i) { auto l = a[s * m + i]; auto r = a[s * m + i + m / 2] * w; a[s * m + i] = l + r; a[s * m + i + m / 2] = l - r; w *= omega; } } } } template <typename T> std::vector<double> convolution(const std::vector<T>& a, const std::vector<T>& b) { const int size = a.size() + b.size() - 1; const int n = std::bit_ceil((unsigned int)size); std::vector<std::complex<double>> na(a.begin(), a.end()), nb(b.begin(), b.end()); na.resize(n); nb.resize(n); fft(na); fft(nb); for (int i = 0; i < n; ++i) na[i] *= nb[i]; ifft(na); std::vector<double> ret(size); for (int i = 0; i < size; ++i) ret[i] = na[i].real() / n; return ret; } #line 5 "test/yosupo/frequency_table_of_tree_distance.test.cpp" #include <bits/stdc++.h> using namespace std; using ll = long long; #define rep(i, s, t) for (int i = (int)(s); i < (int)(t); ++i) #define revrep(i, t, s) for (int i = (int)(t)-1; i >= (int)(s); --i) #define all(x) begin(x), end(x) template <typename T> bool chmax(T& a, const T& b) { return a < b ? (a = b, 1) : 0; } template <typename T> bool chmin(T& a, const T& b) { return a > b ? (a = b, 1) : 0; } int main() { ios_base::sync_with_stdio(false); cin.tie(nullptr); cout << fixed << setprecision(15); int N; cin >> N; vector<vector<int>> G(N); rep(i,0,N-1) { int a, b; cin >> a >> b; G[a].push_back(b); G[b].push_back(a); } vector<int> level, sz, par; tie(level, sz, par) = centroid_decomposition(G); vector<ll> ans(N); vector<double> cnt; auto dfs = [&](auto& dfs, int v, int p, int d, int l) -> void { if (cnt.size() <= d) cnt.push_back(0); ++cnt[d]; for (int c : G[v]) { if (c != p && level[c] > l) { dfs(dfs, c, v, d+1, l); } } }; auto calc = [&](int s) { vector<double> res(sz[s]); cnt.assign(1, 0); for (int c : G[s]) { if (level[c] > level[s]) { dfs(dfs, c, s, 1, level[s]); } } auto conv = convolution(cnt, cnt); rep(d,0,cnt.size()) { res[d] += cnt[d]; } rep(d,0,min(conv.size(), res.size())) { res[d] += conv[d]/2; } for (int c : G[s]) { if (level[c] > level[s]) { cnt.assign(1, 0); dfs(dfs, c, s, 1, level[s]); auto conv = convolution(cnt, cnt); rep(d,0,min(conv.size(), res.size())) { res[d] -= conv[d]/2; } } } rep(d,0,sz[s]) { ans[d] += (ll) round(res[d]); } }; rep(i,0,N) calc(i); rep(d,1,N) cout << ans[d] << (d < N-1 ? " " : "\n"); }