refactor centroid decomposition

hitonanode · hitonanode · commit 5568c4b2f7c1 · 2024-09-25T00:38:47.000+09:00
diff --git a/tree/centroid_decomposition.hpp b/tree/centroid_decomposition.hpp
@@ -1,101 +1,94 @@
 #pragma once
-#include <tuple>
+#include <cassert>
 #include <utility>
 #include <vector>
 
-// CUT begin
-/*
-(Recursive) Centroid Decomposition
-Verification: Codeforces #190 Div.1 C https://codeforces.com/contest/321/submission/59093583
-
-fix_root(int r): Build information of the tree which `r` belongs to.
-detect_centroid(int r): Enumerate centroid(s) of the tree which `r` belongs to.
-*/
+// Centroid Decomposition
+// Verification: https://yukicoder.me/problems/no/2892
+// find_current_centroids(int r, int conn_size): Enumerate centroid(s) of the subtree which `r` belongs to.
 struct CentroidDecomposition {
-    int NO_PARENT = -1;
     int V;
-    int E;
-    std::vector<std::vector<std::pair<int, int>>> to; // (node_id, edge_id)
-    std::vector<int> par;                             // parent node_id par[root] = -1
-    std::vector<std::vector<int>> chi;                // children id's
-    std::vector<int> subtree_size;                    // size of each subtree
-    std::vector<int> available_edge;                  // If 0, ignore the corresponding edge.
-
-    CentroidDecomposition(int v = 0)
-        : V(v), E(0), to(v), par(v, NO_PARENT), chi(v), subtree_size(v) {}
-    CentroidDecomposition(const std::vector<std::vector<int>> &to_)
-        : CentroidDecomposition(to_.size()) {
-        for (int i = 0; i < V; i++) {
-            for (auto j : to_[i]) {
-                if (i < j) { add_edge(i, j); }
-            }
+    std::vector<std::vector<int>> to;
+
+private:
+    std::vector<int> is_alive;
+    std::vector<int> subtree_size;
+
+    void decompose(int r, int conn_size, auto callback) {
+
+        const int c = find_current_centroids(r, conn_size).first;
+        is_alive.at(c) = 0;
+
+        callback(c);
+
+        for (int nxt : to.at(c)) {
+            if (!is_alive.at(nxt)) continue;
+            int next_size = subtree_size.at(nxt);
+            if (subtree_size.at(nxt) > subtree_size.at(c))
+                next_size = subtree_size.at(r) - subtree_size.at(c);
+            decompose(nxt, next_size, callback);
         }
     }
 
+public:
+    CentroidDecomposition(int v = 0) : V(v), to(v), is_alive(v, 1), subtree_size(v) {}
+
+    CentroidDecomposition(int v, const std::vector<std::pair<int, int>> &tree_edges)
+        : CentroidDecomposition(v) {
+        for (auto e : tree_edges) add_edge(e.first, e.second);
+    }
+
     void add_edge(int v1, int v2) {
-        to[v1].emplace_back(v2, E), to[v2].emplace_back(v1, E), E++;
-        available_edge.emplace_back(1);
+        assert(0 <= v1 and v1 < V and 0 <= v2 and v2 < V);
+        assert(v1 != v2);
+        to.at(v1).push_back(v2), to.at(v2).emplace_back(v1);
     }
 
-    int _dfs_fixroot(int now, int prv) {
-        chi[now].clear(), subtree_size[now] = 1;
-        for (auto nxt : to[now]) {
-            if (nxt.first != prv and available_edge[nxt.second]) {
-                par[nxt.first] = now, chi[now].push_back(nxt.first);
-                subtree_size[now] += _dfs_fixroot(nxt.first, now);
+    std::pair<int, int> find_current_centroids(int r, int conn_size) {
+        assert(is_alive.at(r));
+
+        const int thres = conn_size / 2;
+
+        int c1 = -1, c2 = -1;
+
+        auto rec_search = [&](auto &&self, int now, int prv) -> void {
+            bool is_centroid = true;
+            subtree_size.at(now) = 1;
+            for (int nxt : to.at(now)) {
+                if (nxt == prv or !is_alive.at(nxt)) continue;
+                self(self, nxt, now);
+                subtree_size.at(now) += subtree_size.at(nxt);
+                if (subtree_size.at(nxt) > thres) is_centroid = false;
             }
-        }
-        return subtree_size[now];
-    }
+            if (conn_size - subtree_size.at(now) > thres) is_centroid = false;
 
-    void fix_root(int root) {
-        par[root] = NO_PARENT;
-        _dfs_fixroot(root, -1);
+            if (is_centroid) (c1 < 0 ? c1 : c2) = now;
+        };
+        rec_search(rec_search, r, -1);
+
+        return {c1, c2};
     }
 
-    //// Centroid Decpmposition ////
-    std::vector<int> centroid_cand_tmp;
-    void _dfs_detect_centroids(int now, int prv, int n) {
-        bool is_centroid = true;
-        for (auto nxt : to[now]) {
-            if (nxt.first != prv and available_edge[nxt.second]) {
-                _dfs_detect_centroids(nxt.first, now, n);
-                if (subtree_size[nxt.first] > n / 2) is_centroid = false;
+    void run(int r, auto callback) {
+        int conn_size = 0;
+
+        auto rec = [&](auto &&self, int now, int prv) -> void {
+            ++conn_size;
+            is_alive.at(now) = 1;
+
+            for (int nxt : to.at(now)) {
+                if (nxt == prv) continue;
+                self(self, nxt, now);
             }
-        }
-        if (n - subtree_size[now] > n / 2) is_centroid = false;
-        if (is_centroid) centroid_cand_tmp.push_back(now);
-    }
-    std::pair<int, int> detect_centroids(int r) { // ([centroid_node_id1], ([centroid_node_id2]|-1))
-        centroid_cand_tmp.clear();
-        while (par[r] != NO_PARENT) r = par[r];
-        int n = subtree_size[r];
-        _dfs_detect_centroids(r, -1, n);
-        if (centroid_cand_tmp.size() == 1)
-            return std::make_pair(centroid_cand_tmp[0], -1);
-        else
-            return std::make_pair(centroid_cand_tmp[0], centroid_cand_tmp[1]);
-    }
+        };
+        rec(rec, r, -1);
 
-    std::vector<int> _cd_vertices;
-    void _centroid_decomposition(int now) {
-        fix_root(now);
-        now = detect_centroids(now).first;
-        _cd_vertices.emplace_back(now);
-        /*
-        do something
-        */
-        for (auto p : to[now]) {
-            int nxt, eid;
-            std::tie(nxt, eid) = p;
-            if (available_edge[eid] == 0) continue;
-            available_edge[eid] = 0;
-            _centroid_decomposition(nxt);
-        }
+        decompose(r, conn_size, callback);
     }
-    std::vector<int> centroid_decomposition(int x) {
-        _cd_vertices.clear();
-        _centroid_decomposition(x);
-        return _cd_vertices;
+
+    std::vector<int> centroid_decomposition(int r) {
+        std::vector<int> res;
+        run(r, [&](int v) { res.push_back(v); });
+        return res;
     }
 };
diff --git a/tree/centroid_decomposition.md b/tree/centroid_decomposition.md
@@ -0,0 +1,21 @@
+---
+title: Centroid decomposition （森の重心分解）
+documentation_of: ./centroid_decomposition.hpp
+---
+
+与えられた森について，指定された頂点に関する連結成分の重心分解を行う．
+
+## 使用方法
+
+```cpp
+int v = 0;
+
+// 頂点 v の連結成分を重心分解していく
+for (int c : cd.centroid_decomposition(v)) {
+    // 頂点 c を削除する
+}
+```
+
+## 問題例
+
+- [No.2892 Lime and Karin - yukicoder](https://yukicoder.me/problems/no/2892)
