nit

Author: lrvideckis

library/trees/centroid_decomp.hpp

@@ -1,12 +1,12 @@
 #pragma once
 //! @code
 //!   vector<basic_string<int>> g(n);
-//!   centroid(g, [&](int cent, int par_cent) {});
+//!   vi par = cd(g, [&](int cent) {});
 //! @endcode
 //! @time O(n log n)
 //! @space O(n)
-void centroid(auto& g, auto f) {
-  vi siz(sz(g));
+vi cd(auto& g, auto f) {
+  vi p(sz(g), -1), siz(p);
   auto ctd = [&](auto ctd, int u, int p, int n) -> int {
     siz[u] = 1;
     for (int v : g[u])
@@ -16,10 +16,10 @@ void centroid(auto& g, auto f) {
       }
     return 2 * siz[u] >= n ? siz[p] = n - siz[u], u : -1;
   };
-  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]);
+  auto dfs = [&](auto dfs, int u, int n) -> int {
+    f(u = ctd(ctd, u, u, n));
+    for (int v : g[u]) erase(g[v], u), p[dfs(dfs, v, siz[v])] = u;
+    return u;
   };
-  dfs(dfs, 0, -1, sz(g));
+  return dfs(dfs, 0, sz(g)), p;
 }

tests/library_checker_aizu_tests/cd_asserts.hpp

@@ -3,25 +3,24 @@
 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);
+  vi par = cd(adj, [&](int cent) -> void {
     assert(decomp_size[cent] == -1);
     auto dfs = [&](auto&& self, int u, int p) -> int {
-      naive_par_decomp[u] = cent;
+      if (p != -1) 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]);
+      assert(1 <= sz_subtree && 2 * sz_subtree <= decomp_size[cent]);
     }
   });
-  rep(i, 0, sz(adj)) assert(decomp_size[i] >= 1);
+  assert(par == naive_par_decomp);
+  rep(i, 0, sz(adj)) {
+    assert(decomp_size[i] >= 1);
+    if (par[i] != -1) assert(2 * decomp_size[i] <= decomp_size[par[i]]);
+  }
 }

tests/library_checker_aizu_tests/trees/cd_count_paths_per_length.test.cpp

@@ -11,7 +11,7 @@
 //! 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) {
+  cd(adj, [&](int cent) {
     vector<vector<double>> child_depths;
     for (int v : adj[cent]) {
       child_depths.emplace_back(1, 0.0);
@@ -29,15 +29,13 @@ vector<ll> count_paths_per_length(vector<vi> adj) {
         swap(q, new_q);
       }
     }
-    sort(all(child_depths),
-      [&](auto& x, auto& y) { return sz(x) < sz(y); });
+    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];
+      rep(i, 1, sz(cnt_depth)) total_depth[i] += cnt_depth[i];
     }
   });
   return num_paths;