Add HLD and LCA

Author: adamant-pwn

cp-algo/graph/base.hpp

@@ -23,16 +23,17 @@ namespace cp_algo::graph {
             }
             return gt;
         }
-        void add_edge(node_index u, edge_t e) {
+        edge_index add_edge(node_index u, edge_t e) {
             edge_index idx = (edge_index)size(E);
             E.push_back(e);
             adj.push(u, idx);
             if constexpr (mode == undirected) {
                 adj.push(e.traverse(u), idx);
             }
+            return idx;
         }
-        void add_edge(node_index u, auto... Args) {
-            add_edge(u, edge_t(u, Args...));
+        edge_index add_edge(node_index u, auto... Args) {
+            return add_edge(u, edge_t(u, Args...));
         }
         void read_edges(node_index m) {
             adj.reserve(mode == undirected ? 2 * m : m);

cp-algo/graph/ascending_dfs.hpp renamed to cp-algo/tree/ascending_dfs.hpp

@@ -1,7 +1,7 @@
-#ifndef CP_ALGO_GRAPH_ASCENDING_DFS_HPP
-#define CP_ALGO_GRAPH_ASCENDING_DFS_HPP
+#ifndef CP_ALGO_TREE_ASCENDING_DFS_HPP
+#define CP_ALGO_TREE_ASCENDING_DFS_HPP
 
-#include "base.hpp"
+#include "../graph/base.hpp"
 #include <cassert>
 #include <vector>
 #include <ranges>
@@ -65,4 +65,4 @@ namespace cp_algo::graph {
         }, callback, root);
     }
 }
-#endif // CP_ALGO_GRAPH_ASCENDING_DFS_HPP
+#endif // CP_ALGO_TREE_ASCENDING_DFS_HPP

cp-algo/graph/tree_diameter.hpp renamed to cp-algo/tree/diameter.hpp

@@ -1,7 +1,7 @@
-#ifndef CP_ALGO_GRAPH_TREE_DIAMETER_HPP
-#define CP_ALGO_GRAPH_TREE_DIAMETER_HPP
+#ifndef CP_ALGO_TREE_DIAMETER_HPP
+#define CP_ALGO_TREE_DIAMETER_HPP
 #include "ascending_dfs.hpp"
-#include "base.hpp"
+#include "../graph/base.hpp"
 #include <tuple>
 #include <string>
 #include <algorithm>
@@ -59,4 +59,4 @@ namespace cp_algo::graph {
         }
     }
 }
-#endif // CP_ALGO_GRAPH_TREE_DIAMETER_HPP
+#endif // CP_ALGO_TREE_DIAMETER_HPP

cp-algo/tree/hld.hpp

@@ -0,0 +1,68 @@
+#ifndef CP_ALGO_TREE_HLD_HPP
+#define CP_ALGO_TREE_HLD_HPP
+#include "ascending_dfs.hpp"
+#include "../graph/base.hpp"
+#include "../util/big_alloc.hpp"
+#include <vector>
+#include <ranges>
+
+namespace cp_algo::graph {
+    struct heavy_light {
+        big_vector<node_index> size, in, up, par;
+        
+        template<undirected_graph_type graph>
+        heavy_light(graph const& g, node_index root = 0, std::vector<edge_index> const* parents_ptr = nullptr):
+            size(g.n(), 1), in(g.n()), up(g.n()), par(g.n()) {
+            big_vector<node_index> topsort;
+            topsort.reserve(g.n());
+            auto push_size = [&](node_index v, edge_index e) {
+                topsort.push_back(v);
+                if (v != root) {
+                    auto p = g.edge(e).traverse(v);
+                    size[p] += size[v];
+                    par[v] = p;
+                }
+            };
+            if (parents_ptr) {
+                parent_dfs(g, *parents_ptr, push_size);
+            } else {
+                xor_dfs(g, push_size, root);
+            }
+            par[root] = up[root] = root;
+            for(auto v: topsort | std::views::reverse) {
+                node_index big = -1;
+                for(auto e: g.outgoing(v)) {
+                    auto u = g.edge(e).traverse(v);
+                    if (size[u] > size[v]) continue;
+                    if (big == -1 || size[u] > size[big]) {
+                        big = u;
+                    }
+                }
+                int t = in[v] + (big == -1 ? 0 : size[big]);
+                for(auto e: g.outgoing(v)) {
+                    auto u = g.edge(e).traverse(v);
+                    if (size[u] > size[v]) continue;
+                    if (u == big) {
+                        in[u] = in[v] + 1;
+                        up[u] = up[v];
+                    } else {
+                        in[u] = t + 1;
+                        t += size[u];
+                        up[u] = u;
+                    }
+                }
+            }
+        }
+        node_index lca(node_index a, node_index b) {
+            while (up[a] != up[b]) {
+                if (in[up[a]] < in[up[b]]) {
+                    b = par[up[b]];
+                } else {
+                    a = par[up[a]];
+                }
+            }
+            return in[a] < in[b] ? a : b;
+        }
+    };
+}
+#endif // CP_ALGO_TREE_HLD_HPP

verify/graph/minimum_diameter_spanning_tree.test.cpp

@@ -4,7 +4,7 @@
 #include <iostream>
 #include "blazingio/blazingio.min.hpp"
 #include "cp-algo/graph/shortest_path.hpp"
-#include "cp-algo/graph/tree_diameter.hpp"
+#include "cp-algo/tree/diameter.hpp"
 #include <bits/stdc++.h>
 
 using namespace std;

verify/graph/tree_diameter.test.cpp renamed to verify/tree/diameter.test.cpp

@@ -3,7 +3,7 @@
 #pragma GCC optimize("Ofast,unroll-loops")
 #include <iostream>
 #include "blazingio/blazingio.min.hpp"
-#include "cp-algo/graph/tree_diameter.hpp"
+#include "cp-algo/tree/diameter.hpp"
 #include <bits/stdc++.h>
 
 using namespace std;

verify/tree/lca.test.cpp

@@ -0,0 +1,40 @@
+// @brief Lowest Common Ancestor
+#define PROBLEM "https://judge.yosupo.jp/problem/lca"
+#pragma GCC optimize("Ofast,unroll-loops")
+#include <iostream>
+#include "blazingio/blazingio.min.hpp"
+#include "cp-algo/tree/hld.hpp"
+#include <bits/stdc++.h>
+
+using namespace std;
+using namespace cp_algo::graph;
+
+void solve() {
+    int n, q;
+    cin >> n >> q;
+    graph g(n);
+    vector<edge_index> p(n);
+    p[0] = -1;
+    for(int i = 1; i < n; i++) {
+        node_index v;
+        cin >> v;
+        p[i] = g.add_edge(v, i);
+    }
+    heavy_light hld(g, 0, &p);
+    for(int i = 0; i < q; i++) {
+        node_index u, v;
+        cin >> u >> v;
+        cout << hld.lca(u, v) << '\n';
+    }
+}
+
+signed main() {
+    //freopen("input.txt", "r", stdin);
+    ios::sync_with_stdio(0);
+    cin.tie(0);
+    int t = 1;
+    //cin >> t;
+    while(t--) {
+        solve();
+    }
+}