add back centroid test

Author: lrvideckis

tests/library_checker_aizu_tests/cd_asserts.hpp

@@ -1 +1,27 @@
 #pragma once
+#include "../../library/trees/centroid_decomp.hpp"
+void cd_asserts(vector<vector<int>> adj) {
+  vector<int> decomp_size(sz(adj), -1);
+  vector<int> naive_par_decomp(sz(adj), -1);
+  centroid(adj, [&](int cent, int par_cent) -> void {
+    assert(naive_par_decomp[cent] == par_cent);
+    assert(decomp_size[cent] == -1);
+    auto dfs = [&](auto&& self, int u, int p) -> int {
+      naive_par_decomp[u] = cent;
+      int sub_size = 1;
+      for (int v : adj[u])
+        if (v != p) sub_size += self(self, v, u);
+      return sub_size;
+    };
+    decomp_size[cent] = dfs(dfs, cent, -1);
+    if (par_cent != -1)
+      assert(1 <= decomp_size[cent] &&
+        2 * decomp_size[cent] <= decomp_size[par_cent]);
+    for (int u : adj[cent]) {
+      int sz_subtree = dfs(dfs, u, cent);
+      assert(1 <= sz_subtree &&
+        2 * sz_subtree <= decomp_size[cent]);
+    }
+  });
+  rep(i, 0, sz(adj)) assert(decomp_size[i] >= 1);
+}

tests/library_checker_aizu_tests/handmade_tests/count_paths.test.cpp

@@ -6,31 +6,7 @@
 #include "../../../kactl/content/numerical/FastFourierTransform.h"
 #include "../../../library/trees/edge_cd.hpp"
 #include "../../../library/trees/centroid_decomp.hpp"
-void cd_asserts(vector<vector<int>> adj) {
-  vector<int> decomp_size(sz(adj), -1);
-  vector<int> naive_par_decomp(sz(adj), -1);
-  centroid(adj, [&](int cent, int par_cent) -> void {
-    assert(naive_par_decomp[cent] == par_cent);
-    assert(decomp_size[cent] == -1);
-    auto dfs = [&](auto&& self, int u, int p) -> int {
-      naive_par_decomp[u] = cent;
-      int sub_size = 1;
-      for (int v : adj[u])
-        if (v != p) sub_size += self(self, v, u);
-      return sub_size;
-    };
-    decomp_size[cent] = dfs(dfs, cent, -1);
-    if (par_cent != -1)
-      assert(1 <= decomp_size[cent] &&
-        2 * decomp_size[cent] <= decomp_size[par_cent]);
-    for (int u : adj[cent]) {
-      int sz_subtree = dfs(dfs, u, cent);
-      assert(1 <= sz_subtree &&
-        2 * sz_subtree <= decomp_size[cent]);
-    }
-  });
-  rep(i, 0, sz(adj)) assert(decomp_size[i] >= 1);
-}
+#include "../cd_asserts.hpp"
 //! @param adj unrooted, connected forest
 //! @param k number of edges
 //! @returns array `num_paths` where `num_paths[i]` =

tests/library_checker_aizu_tests/trees/cd_count_paths_per_length.test.cpp

@@ -0,0 +1,61 @@
+#define PROBLEM \
+  "https://judge.yosupo.jp/problem/frequency_table_of_tree_distance"
+#include "../template.hpp"
+#include "../cd_asserts.hpp"
+#include "../../../kactl/content/numerical/FastFourierTransform.h"
+//! @param adj unrooted, connected forest
+//! @returns array `num_paths` where `num_paths[i]` = # of
+//! paths in tree with `i` edges. `num_paths[1]` = # edges
+//! @time O(n log^2 n)
+//! @space this function allocates/returns various vectors
+//! which are each O(n)
+vector<ll> count_paths_per_length(vector<vi> adj) {
+  vector<ll> num_paths(sz(adj));
+  centroid(adj, [&](int cent, int) {
+    vector<vector<double>> child_depths;
+    for (int v : adj[cent]) {
+      child_depths.emplace_back(1, 0.0);
+      for (queue<pii> q({{v, cent}}); !empty(q);) {
+        child_depths.back().push_back(sz(q));
+        queue<pii> new_q;
+        while (!empty(q)) {
+          auto [u, p] = q.front();
+          q.pop();
+          for (int w : adj[u]) {
+            if (w == p) continue;
+            new_q.push({w, u});
+          }
+        }
+        swap(q, new_q);
+      }
+    }
+    sort(all(child_depths),
+      [&](auto& x, auto& y) { return sz(x) < sz(y); });
+    vector total_depth(1, 1.0);
+    for (auto& cnt_depth : child_depths) {
+      auto prod = conv(total_depth, cnt_depth);
+      rep(i, 1, sz(prod)) num_paths[i] += llround(prod[i]);
+      total_depth.resize(sz(cnt_depth));
+      rep(i, 1, sz(cnt_depth)) total_depth[i] +=
+        cnt_depth[i];
+    }
+  });
+  return num_paths;
+}
+int main() {
+  cin.tie(0)->sync_with_stdio(0);
+  int n;
+  cin >> n;
+  vector<vector<int>> adj(n);
+  for (int i = 0; i < n - 1; i++) {
+    int u, v;
+    cin >> u >> v;
+    adj[u].push_back(v);
+    adj[v].push_back(u);
+  }
+  cd_asserts(adj);
+  vector<ll> cnt_len = count_paths_per_length(adj);
+  for (int i = 1; i < n; i++) cout << cnt_len[i] << " ";
+  cout << '\n';
+  return 0;
+}