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: 2024-09-25 00:42:32+09:00
- Problem: https://judge.yosupo.jp/problem/frequency_table_of_tree_distance
Depends on
- convolution/fft_double.hpp
- Centroid decomposition （森の重心分解）
(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 <cassert>
#include <utility>
#include <vector>
// Centroid Decomposition
// Verification: https://yukicoder.me/problems/no/2892
// find_current_centroids(int r, int conn_size): Enumerate centroid(s) of the subtree which `r` belongs to.
struct CentroidDecomposition {
int V;
std::vector<std::vector<int>> to;
private:
std::vector<int> is_alive;
std::vector<int> subtree_size;
template <class F> void decompose(int r, int conn_size, F callback) {
const int c = find_current_centroids(r, conn_size).first;
is_alive.at(c) = 0;
callback(c);
for (int nxt : to.at(c)) {
if (!is_alive.at(nxt)) continue;
int next_size = subtree_size.at(nxt);
if (subtree_size.at(nxt) > subtree_size.at(c))
next_size = subtree_size.at(r) - subtree_size.at(c);
decompose(nxt, next_size, callback);
}
}
public:
CentroidDecomposition(int v = 0) : V(v), to(v), is_alive(v, 1), subtree_size(v) {}
CentroidDecomposition(int v, const std::vector<std::pair<int, int>> &tree_edges)
: CentroidDecomposition(v) {
for (auto e : tree_edges) add_edge(e.first, e.second);
}
void add_edge(int v1, int v2) {
assert(0 <= v1 and v1 < V and 0 <= v2 and v2 < V);
assert(v1 != v2);
to.at(v1).push_back(v2), to.at(v2).emplace_back(v1);
}
std::pair<int, int> find_current_centroids(int r, int conn_size) {
assert(is_alive.at(r));
const int thres = conn_size / 2;
int c1 = -1, c2 = -1;
auto rec_search = [&](auto &&self, int now, int prv) -> void {
bool is_centroid = true;
subtree_size.at(now) = 1;
for (int nxt : to.at(now)) {
if (nxt == prv or !is_alive.at(nxt)) continue;
self(self, nxt, now);
subtree_size.at(now) += subtree_size.at(nxt);
if (subtree_size.at(nxt) > thres) is_centroid = false;
}
if (conn_size - subtree_size.at(now) > thres) is_centroid = false;
if (is_centroid) (c1 < 0 ? c1 : c2) = now;
};
rec_search(rec_search, r, -1);
return {c1, c2};
}
template <class F> void run(int r, F callback) {
int conn_size = 0;
auto rec = [&](auto &&self, int now, int prv) -> void {
++conn_size;
is_alive.at(now) = 1;
for (int nxt : to.at(now)) {
if (nxt == prv) continue;
self(self, nxt, now);
}
};
rec(rec, r, -1);
decompose(r, conn_size, callback);
}
std::vector<int> centroid_decomposition(int r) {
std::vector<int> res;
run(r, [&](int v) { res.push_back(v); });
return res;
}
};
#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;
CentroidDecomposition c(to.size());
for (int i = 0; i < int(to.size()); i++) {
for (int j : to[i]) {
if (i < j) c.add_edge(i, j);
}
}
cd = c.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"
// 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';
}