cplib-cpp
This documentation is automatically generated by online-judge-tools/verification-helper
tree/test/frequency_table_of_tree_distance.test.cpp
- View this file on GitHub
- Last update: 2022-01-08 20:23:44+09:00
- Problem: https://judge.yosupo.jp/problem/frequency_table_of_tree_distance
Depends on
- convolution/fft_double.hpp
- tree/centroid_decomposition.hpp
- Frequency table of tree distance
(tree/frequency_table_of_tree_distance.hpp)
Code
#include "../frequency_table_of_tree_distance.hpp"
#include "../../convolution/fft_double.hpp"
#include <algorithm>
#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);
}
auto ret =
frequency_table_of_tree_distance(to).solve<long long, fftconv>(std::vector<long long>(N, 1));
for (int i = 1; i < N; i++) cout << ret[i] << ' ';
cout << '\n';
}#line 2 "tree/centroid_decomposition.hpp"
#include <tuple>
#include <utility>
#include <vector>
// 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 2 "tree/frequency_table_of_tree_distance.hpp"
#include <algorithm>
#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 2 "convolution/fft_double.hpp"
#include <complex>
#line 5 "convolution/fft_double.hpp"
// CUT begin
// Convolution by FFT (Fast Fourier Transform)
// Algorithm based on http://kirika-comp.hatenablog.com/entry/2018/03/12/210446
// Verified: ATC001C (168 ms) https://atcoder.jp/contests/atc001/submissions/9243440
using cmplx = std::complex<double>;
void fft(int N, std::vector<cmplx> &a, double dir) {
int i = 0;
for (int j = 1; j < N - 1; j++) {
for (int k = N >> 1; k > (i ^= k); k >>= 1)
;
if (j < i) std::swap(a[i], a[j]);
}
std::vector<cmplx> zeta_pow(N);
for (int i = 0; i < N; i++) {
double theta = M_PI / N * i * dir;
zeta_pow[i] = {cos(theta), sin(theta)};
}
for (int m = 1; m < N; m *= 2) {
for (int y = 0; y < m; y++) {
cmplx fac = zeta_pow[N / m * y];
for (int x = 0; x < N; x += 2 * m) {
int u = x + y;
int v = x + y + m;
cmplx s = a[u] + fac * a[v];
cmplx t = a[u] - fac * a[v];
a[u] = s;
a[v] = t;
}
}
}
}
template <typename T>
std::vector<cmplx> conv_cmplx(const std::vector<T> &a, const std::vector<T> &b) {
int N = 1;
while (N < (int)a.size() + (int)b.size()) N *= 2;
std::vector<cmplx> a_(N), b_(N);
for (int i = 0; i < (int)a.size(); i++) a_[i] = a[i];
for (int i = 0; i < (int)b.size(); i++) b_[i] = b[i];
fft(N, a_, 1);
fft(N, b_, 1);
for (int i = 0; i < N; i++) a_[i] *= b_[i];
fft(N, a_, -1);
for (int i = 0; i < N; i++) a_[i] /= N;
return a_;
}
// retval[i] = \sum_j a[j]b[i - j]
// Requirement: length * max(a) * max(b) < 10^15
template <typename T>
std::vector<long long int> fftconv(const std::vector<T> &a, const std::vector<T> &b) {
std::vector<cmplx> ans = conv_cmplx(a, b);
std::vector<long long int> ret(ans.size());
for (int i = 0; i < (int)ans.size(); i++) ret[i] = floor(ans[i].real() + 0.5);
ret.resize(a.size() + b.size() - 1);
return ret;
}
#line 4 "tree/test/frequency_table_of_tree_distance.test.cpp"
#include <iostream>
#line 6 "tree/test/frequency_table_of_tree_distance.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);
}
auto ret =
frequency_table_of_tree_distance(to).solve<long long, fftconv>(std::vector<long long>(N, 1));
for (int i = 1; i < N; i++) cout << ret[i] << ' ';
cout << '\n';
}