|
3 | 3 | #include "../template.hpp" |
4 | 4 | #include "../../../library/contest/random.hpp" |
5 | 5 | #include "../../../library/trees/lca_rmq.hpp" |
6 | | -#include "../cd_asserts.hpp" |
7 | 6 | #include "../../../kactl/content/numerical/FastFourierTransform.h" |
8 | 7 | #include "../../../library/trees/edge_cd.hpp" |
| 8 | +#include "../../../library/trees/centroid_decomp.hpp" |
| 9 | +void cd_asserts(vector<vector<int>> adj) { |
| 10 | + vector<int> decomp_size(sz(adj), -1); |
| 11 | + vector<int> naive_par_decomp(sz(adj), -1); |
| 12 | + centroid(adj, [&](int cent, int par_cent) -> void { |
| 13 | + assert(naive_par_decomp[cent] == par_cent); |
| 14 | + assert(decomp_size[cent] == -1); |
| 15 | + auto dfs = [&](auto&& self, int u, int p) -> int { |
| 16 | + naive_par_decomp[u] = cent; |
| 17 | + int sub_size = 1; |
| 18 | + for (int v : adj[u]) |
| 19 | + if (v != p) sub_size += self(self, v, u); |
| 20 | + return sub_size; |
| 21 | + }; |
| 22 | + decomp_size[cent] = dfs(dfs, cent, -1); |
| 23 | + if (par_cent != -1) |
| 24 | + assert(1 <= decomp_size[cent] && |
| 25 | + 2 * decomp_size[cent] <= decomp_size[par_cent]); |
| 26 | + for (int u : adj[cent]) { |
| 27 | + int sz_subtree = dfs(dfs, u, cent); |
| 28 | + assert(1 <= sz_subtree && |
| 29 | + 2 * sz_subtree <= decomp_size[cent]); |
| 30 | + } |
| 31 | + }); |
| 32 | + rep(i, 0, sz(adj)) assert(decomp_size[i] >= 1); |
| 33 | +} |
9 | 34 | //! @param adj unrooted, connected forest |
10 | 35 | //! @param k number of edges |
11 | 36 | //! @returns array `num_paths` where `num_paths[i]` = |
|
14 | 39 | //! @time O(n log n) |
15 | 40 | //! @space this function allocates/returns various vectors |
16 | 41 | //! which are all O(n) |
17 | | -vector<ll> count_paths_per_node(const vector<vi>& adj, |
18 | | - int k) { |
| 42 | +vector<ll> count_paths_per_node(vector<vi> adj, int k) { |
19 | 43 | 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 | | - }); |
| 44 | + centroid(adj, [&](int cent, int) { |
| 45 | + vector pre_d{1}, cur_d{0}; |
| 46 | + auto dfs = [&](auto&& self, int u, int p, |
| 47 | + int d) -> ll { |
| 48 | + if (d > k) return 0LL; |
| 49 | + if (sz(cur_d) <= d) cur_d.push_back(0); |
| 50 | + cur_d[d]++; |
| 51 | + ll cnt = 0; |
| 52 | + if (k - d < sz(pre_d)) cnt += pre_d[k - d]; |
| 53 | + for (int c : adj[u]) |
| 54 | + if (c != p) cnt += self(self, c, u, d + 1); |
| 55 | + num_paths[u] += cnt; |
| 56 | + return cnt; |
| 57 | + }; |
| 58 | + auto dfs_child = [&](int child) -> ll { |
| 59 | + ll cnt = dfs(dfs, child, cent, 1); |
| 60 | + pre_d.resize(sz(cur_d)); |
| 61 | + for (int i = 1; i < sz(cur_d) && cur_d[i]; i++) |
| 62 | + pre_d[i] += cur_d[i], cur_d[i] = 0; |
| 63 | + return cnt; |
| 64 | + }; |
| 65 | + for (int child : adj[cent]) |
| 66 | + num_paths[cent] += dfs_child(child); |
| 67 | + pre_d = cur_d = {0}; |
| 68 | + for (int child : adj[cent] | views::reverse) |
| 69 | + dfs_child(child); |
| 70 | + }); |
48 | 71 | return num_paths; |
49 | 72 | } |
50 | 73 | vector<vector<ll>> naive(const vector<vi>& adj) { |
|
0 commit comments