Skip to content

Commit 0448507

Browse files
committed
another fix
1 parent 121d8b4 commit 0448507

File tree

1 file changed

+28
-1
lines changed

1 file changed

+28
-1
lines changed

tests/library_checker_aizu_tests/trees/edge_cd_count_paths_per_length.test.cpp

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,34 @@
22
"https://judge.yosupo.jp/problem/frequency_table_of_tree_distance"
33
#include "../template.hpp"
44
#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+
}
633
int main() {
734
cin.tie(0)->sync_with_stdio(0);
835
int n;

0 commit comments

Comments
 (0)