22 " https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"
33#include " ../template.hpp"
44#include " ../../../library/contest/random.hpp"
5- #include " ../../../library/trees/uncommon/count_paths_per_node.hpp"
65#include " ../../../library/trees/lca_rmq.hpp"
7- #include " ../../../library/trees/uncommon/count_paths_per_length.hpp"
86#include " ../cd_asserts.hpp"
7+ #include " ../../../kactl/content/numerical/FastFourierTransform.h"
8+ #include " ../../../library/trees/edge_cd.hpp"
9+ // ! @param adj unrooted, connected forest
10+ // ! @param k number of edges
11+ // ! @returns array `num_paths` where `num_paths[i]` =
12+ // ! number of paths with k edges where node `i` is on the
13+ // ! path. 0-based nodes.
14+ // ! @time O(n log n)
15+ // ! @space this function allocates/returns various vectors
16+ // ! which are all O(n)
17+ vector<ll> count_paths_per_node (const vector<vi>& adj,
18+ int k) {
19+ vector<ll> num_paths (sz (adj));
20+ centroid (adj,
21+ [&](const vector<vi>& cd_adj, int cent, int ) {
22+ vector pre_d{1 }, cur_d{0 };
23+ auto dfs = [&](auto && self, int u, int p,
24+ int d) -> ll {
25+ if (d > k) return 0LL ;
26+ if (sz (cur_d) <= d) cur_d.push_back (0 );
27+ cur_d[d]++;
28+ ll cnt = 0 ;
29+ if (k - d < sz (pre_d)) cnt += pre_d[k - d];
30+ for (int c : cd_adj[u])
31+ if (c != p) cnt += self (self, c, u, d + 1 );
32+ num_paths[u] += cnt;
33+ return cnt;
34+ };
35+ auto dfs_child = [&](int child) -> ll {
36+ ll cnt = dfs (dfs, child, cent, 1 );
37+ pre_d.resize (sz (cur_d));
38+ for (int i = 1 ; i < sz (cur_d) && cur_d[i]; i++)
39+ pre_d[i] += cur_d[i], cur_d[i] = 0 ;
40+ return cnt;
41+ };
42+ for (int child : cd_adj[cent])
43+ num_paths[cent] += dfs_child (child);
44+ pre_d = cur_d = {0 };
45+ for (int child : cd_adj[cent] | views::reverse)
46+ dfs_child (child);
47+ });
48+ return num_paths;
49+ }
950vector<vector<ll>> naive (const vector<vi>& adj) {
1051 LCA lc (adj);
1152 int n = sz (adj);
@@ -21,6 +62,32 @@ vector<vector<ll>> naive(const vector<vi>& adj) {
2162 }
2263 return cnts_naive;
2364}
65+ // ! @param adj unrooted, connected tree
66+ // ! @returns array `num_paths` where `num_paths[i]` = # of
67+ // ! paths in tree with `i` edges. `num_paths[1]` = # edges
68+ // ! @time O(n * logφ(n) * log2(n))
69+ // ! @space this function allocates/returns various vectors
70+ // ! which are each O(n)
71+ vector<ll> count_paths_per_length (const vector<vi>& adj) {
72+ vector<ll> num_paths (sz (adj));
73+ if (sz (adj) >= 2 ) num_paths[1 ] = sz (adj) - 1 ;
74+ edge_cd (adj,
75+ [&](const vector<vi>& cd_adj, int cent, int split) {
76+ vector<vector<double >> cnt (2 , vector<double >(1 ));
77+ auto dfs = [&](auto && self, int u, int p, int d,
78+ int side) -> void {
79+ if (sz (cnt[side]) == d) cnt[side].push_back (0.0 );
80+ cnt[side][d]++;
81+ for (int c : cd_adj[u])
82+ if (c != p) self (self, c, u, 1 + d, side);
83+ };
84+ rep (i, 0 , sz (cd_adj[cent]))
85+ dfs (dfs, cd_adj[cent][i], cent, 1 , i < split);
86+ vector<double > prod = conv (cnt[0 ], cnt[1 ]);
87+ rep (i, 0 , sz (prod)) num_paths[i] += llround (prod[i]);
88+ });
89+ return num_paths;
90+ }
2491int main () {
2592 cin.tie (0 )->sync_with_stdio (0 );
2693 for (int n = 1 ; n <= 100 ; n++) {
0 commit comments