More golf (#209)

* add back centroid test

* fix

* [auto-verifier] verify commit 31024fe

* golf

* fix

* more golf

* name nit

---------

Co-authored-by: GitHub <noreply@github.com>

Author: lrvideckis

.verify-helper/timestamps.remote.json

@@ -74,7 +74,7 @@
 "tests/library_checker_aizu_tests/graphs/strongly_connected_components_aizu.test.cpp": "2026-04-06 14:41:55 -0600",
 "tests/library_checker_aizu_tests/graphs/strongly_connected_components_lib_checker.test.cpp": "2026-04-06 14:41:55 -0600",
 "tests/library_checker_aizu_tests/graphs/two_edge_components.test.cpp": "2026-04-06 14:41:55 -0600",
-"tests/library_checker_aizu_tests/handmade_tests/count_paths.test.cpp": "2026-04-12 10:51:29 -0600",
+"tests/library_checker_aizu_tests/handmade_tests/count_paths.test.cpp": "2026-04-12 13:18:22 -0600",
 "tests/library_checker_aizu_tests/handmade_tests/dsu.test.cpp": "2026-01-22 10:08:22 -0700",
 "tests/library_checker_aizu_tests/handmade_tests/edge_cd_small_trees.test.cpp": "2026-04-12 10:51:29 -0600",
 "tests/library_checker_aizu_tests/handmade_tests/fib_matrix_expo.test.cpp": "2026-01-28 21:48:16 -0700",
@@ -131,7 +131,7 @@
 "tests/library_checker_aizu_tests/strings/suffix_array_short.test.cpp": "2026-04-12 06:15:25 -0600",
 "tests/library_checker_aizu_tests/strings/trie.test.cpp": "2026-01-17 12:38:18 -0700",
 "tests/library_checker_aizu_tests/strings/wildcard_pattern_matching.test.cpp": "2025-08-05 19:19:23 -0600",
-"tests/library_checker_aizu_tests/trees/count_paths_per_length.test.cpp": "2026-04-12 10:51:29 -0600",
+"tests/library_checker_aizu_tests/trees/cd_count_paths_per_length.test.cpp": "2026-04-12 13:19:58 -0600",
 "tests/library_checker_aizu_tests/trees/edge_cd_contour_range_query.test.cpp": "2026-04-12 10:51:29 -0600",
 "tests/library_checker_aizu_tests/trees/edge_cd_contour_range_update.test.cpp": "2026-04-12 10:51:29 -0600",
 "tests/library_checker_aizu_tests/trees/edge_cd_count_paths_per_length.test.cpp": "2026-04-12 10:51:29 -0600",

library/trees/centroid_decomp.hpp

@@ -7,22 +7,19 @@
 //! @space O(n)
 void centroid(auto& g, auto f) {
   vi siz(sz(g));
-  auto dfs_sz = [&](auto&& dfs_sz, int u, int p) -> void {
+  auto ctd = [&](auto&& ctd, int u, int p, int n) -> int {
     siz[u] = 1;
     for (int v : g[u])
-      if (v != p) dfs_sz(dfs_sz, v, u), siz[u] += siz[v];
+      if (v != p) {
+        if (int c = ctd(ctd, v, u, n); c != -1) return c;
+        siz[u] += siz[v];
+      }
+    return 2 * siz[u] >= n ? siz[p] = n - siz[u], u : -1;
   };
-  auto dfs = [&](auto&& dfs, int u, int p) -> void {
-    dfs_sz(dfs_sz, u, -1);
-    for (int w = -1, sz_root = siz[u];;) {
-      auto big_ch = ranges::find_if(g[u], [&](int v) {
-        return v != w && 2 * siz[v] > sz_root;
-      });
-      if (big_ch == end(g[u])) break;
-      w = u, u = *big_ch;
-    }
-    f(u, p);
-    for (int v : g[u]) erase(g[v], u), dfs(dfs, v, u);
+  auto dfs = [&](auto&& dfs, int u, int p, int n) -> void {
+    f(u = ctd(ctd, u, u, n), p);
+    for (int v : g[u])
+      erase(g[v], u), dfs(dfs, v, u, siz[v]);
   };
-  dfs(dfs, 0, -1);
+  dfs(dfs, 0, -1, sz(g));
 }

library/trees/edge_cd.hpp

@@ -15,21 +15,19 @@
 //! @space O(n)
 template<class G> void edge_cd(vector<G>& g, auto f) {
   vi siz(sz(g));
-  auto cent = [&](auto&& cent, int u, int p,
-                int m) -> int {
+  auto ctd = [&](auto&& ctd, int u, int p, int m) -> int {
     siz[u] = 1;
     for (int v : g[u])
       if (v != p) {
-        int c = cent(cent, v, u, m);
-        if (c != -1) return c;
+        if (int c = ctd(ctd, v, u, m); c != -1) return c;
         siz[u] += siz[v];
       }
     return 2 * siz[u] > m ? siz[p] = m + 1 - siz[u],
                             u : -1;
   };
   auto dfs = [&](auto&& dfs, int u, int m) -> void {
     if (m < 2) return;
-    u = cent(cent, u, u, m);
+    u = ctd(ctd, u, u, m);
     int sum = 0;
     auto it = partition(all(g[u]), [&](int v) {
       ll x = sum + siz[v];

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.emplace(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 (const 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;
+}