|
2 | 2 | "https://judge.yosupo.jp/problem/frequency_table_of_tree_distance" |
3 | 3 | #include "../template.hpp" |
4 | 4 | #include "../edge_cd_asserts.hpp" |
5 | | -#include "../../../library/trees/uncommon/count_paths_per_length.hpp" |
| 5 | +#include "../../../kactl/content/numerical/FastFourierTransform.h" |
| 6 | +#include "../../../library/trees/edge_cd.hpp" |
| 7 | +//! @param adj unrooted, connected tree |
| 8 | +//! @returns array `num_paths` where `num_paths[i]` = # of |
| 9 | +//! paths in tree with `i` edges. `num_paths[1]` = # edges |
| 10 | +//! @time O(n * logφ(n) * log2(n)) |
| 11 | +//! @space this function allocates/returns various vectors |
| 12 | +//! which are each O(n) |
| 13 | +vector<ll> count_paths_per_length(const vector<vi>& adj) { |
| 14 | + vector<ll> num_paths(sz(adj)); |
| 15 | + if (sz(adj) >= 2) num_paths[1] = sz(adj) - 1; |
| 16 | + edge_cd(adj, |
| 17 | + [&](const vector<vi>& cd_adj, int cent, int split) { |
| 18 | + vector<vector<double>> cnt(2, vector<double>(1)); |
| 19 | + auto dfs = [&](auto&& self, int u, int p, int d, |
| 20 | + int side) -> void { |
| 21 | + if (sz(cnt[side]) == d) cnt[side].push_back(0.0); |
| 22 | + cnt[side][d]++; |
| 23 | + for (int c : cd_adj[u]) |
| 24 | + if (c != p) self(self, c, u, 1 + d, side); |
| 25 | + }; |
| 26 | + rep(i, 0, sz(cd_adj[cent])) |
| 27 | + dfs(dfs, cd_adj[cent][i], cent, 1, i < split); |
| 28 | + vector<double> prod = conv(cnt[0], cnt[1]); |
| 29 | + rep(i, 0, sz(prod)) num_paths[i] += llround(prod[i]); |
| 30 | + }); |
| 31 | + return num_paths; |
| 32 | +} |
6 | 33 | int main() { |
7 | 34 | cin.tie(0)->sync_with_stdio(0); |
8 | 35 | int n; |
|
0 commit comments