golf

Author: lrvideckis

library/trees/centroid_decomp.hpp

@@ -7,22 +7,20 @@
 //! @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 cnt = [&](auto&& cnt, 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) {
+        int c = cnt(cnt, v, u, n);
+        if (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 = cnt(cnt, 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,12 +15,11 @@
 //! @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 cnt = [&](auto&& cnt, 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);
+        int c = cnt(cnt, v, u, m);
         if (c != -1) return c;
         siz[u] += siz[v];
       }
@@ -29,7 +28,7 @@ template<class G> void edge_cd(vector<G>& g, auto f) {
   };
   auto dfs = [&](auto&& dfs, int u, int m) -> void {
     if (m < 2) return;
-    u = cent(cent, u, u, m);
+    u = cnt(cnt, u, u, m);
     int sum = 0;
     auto it = partition(all(g[u]), [&](int v) {
       ll x = sum + siz[v];