modernize edge CD

Author: lrvideckis

library/trees/edge_cd.hpp

@@ -13,39 +13,36 @@
 //! handle single-edge-paths separately
 //! @time O(n logφ n)
 //! @space O(n)
-template<class F, class G> struct edge_cd {
-  vector<G> adj;
-  F f;
-  vi siz;
-  edge_cd(const vector<G>& adj, F f):
-    adj(adj), f(f), siz(sz(adj)) {
-    dfs(0, sz(adj) - 1);
-  }
-  int find_cent(int u, int p, int m) {
+template<class G>
+void edge_cd(vector<G>& adj, const auto& f) {
+  vi siz(sz(adj));
+  auto find_cent = [&](auto&& self, int u, int p,
+                     int m) -> int {
     siz[u] = 1;
     for (int v : adj[u])
       if (v != p) {
-        int cent = find_cent(v, u, m);
+        int cent = self(self, v, u, m);
         if (cent != -1) return cent;
         siz[u] += siz[v];
       }
     return 2 * siz[u] > m
            ? p >= 0 && (siz[p] = m + 1 - siz[u]),
            u : -1;
-  }
-  void dfs(int u, int m) {
+  };
+  auto dfs = [&](auto&& self, int u, int m) -> void {
     if (m < 2) return;
-    u = find_cent(u, -1, m);
+    u = find_cent(find_cent, u, -1, m);
     int sum = 0;
     auto it = partition(all(adj[u]), [&](int v) {
       ll x = sum + siz[v];
       return x * x < m * (m - x) ? sum += siz[v], 1 : 0;
     });
-    f(adj, u, it - begin(adj[u]));
+    f(u, it - begin(adj[u]));
     G oth(it, end(adj[u]));
     adj[u].erase(it, end(adj[u]));
-    dfs(u, sum);
+    self(self, u, sum);
     swap(adj[u], oth);
-    dfs(u, m - sum);
-  }
+    self(self, u, m - sum);
+  };
+  dfs(dfs, 0, sz(adj) - 1);
 };

tests/library_checker_aizu_tests/trees/count_paths_per_length.test.cpp

@@ -9,24 +9,23 @@
 //! @time O(n * logφ(n) * log2(n))
 //! @space this function allocates/returns various vectors
 //! which are each O(n)
-vector<ll> count_paths_per_length(const vector<vi>& adj) {
+vector<ll> count_paths_per_length(vector<vi>& adj) {
   vector<ll> num_paths(sz(adj));
   if (sz(adj) >= 2) num_paths[1] = sz(adj) - 1;
-  edge_cd(adj,
-    [&](const vector<vi>& cd_adj, int cent, int split) {
-      vector<vector<double>> cnt(2, vector<double>(1));
-      auto dfs = [&](auto&& self, int u, int p, int d,
-                   int side) -> void {
-        if (sz(cnt[side]) == d) cnt[side].push_back(0.0);
-        cnt[side][d]++;
-        for (int c : cd_adj[u])
-          if (c != p) self(self, c, u, 1 + d, side);
-      };
-      rep(i, 0, sz(cd_adj[cent]))
-        dfs(dfs, cd_adj[cent][i], cent, 1, i < split);
-      vector<double> prod = conv(cnt[0], cnt[1]);
-      rep(i, 0, sz(prod)) num_paths[i] += llround(prod[i]);
-    });
+  edge_cd(adj, [&](int cent, int split) {
+    vector<vector<double>> cnt(2, vector<double>(1));
+    auto dfs = [&](auto&& self, int u, int p, int d,
+                 int side) -> void {
+      if (sz(cnt[side]) == d) cnt[side].push_back(0.0);
+      cnt[side][d]++;
+      for (int c : adj[u])
+        if (c != p) self(self, c, u, 1 + d, side);
+    };
+    rep(i, 0, sz(adj[cent]))
+      dfs(dfs, adj[cent][i], cent, 1, i < split);
+    vector<double> prod = conv(cnt[0], cnt[1]);
+    rep(i, 0, sz(prod)) num_paths[i] += llround(prod[i]);
+  });
   return num_paths;
 }
 int main() {