add hld

Author: lrvideckis

library/data_structures/dsu/dsu.hpp

@@ -1,7 +1,8 @@
 #pragma once
-struct dsu {
+// NOLINTNEXTLINE(readability-identifier-naming)
+struct DSU {
   vi e;
-  dsu(int n): e(n, -1) {}
+  DSU(int n): e(n, -1) {}
   int size(int x) { return -e[go(x)]; }
   int go(int x) { return e[x] < 0 ? x : e[x] = go(e[x]); }
   bool join(int a, int b) {

library/trees/hld.hpp

@@ -0,0 +1,49 @@
+#pragma once
+//! https://github.com/kth-competitive-programming/kactl/blob/main/content/graph/HLD.h
+//! @code
+//!   vector<basic_string<int>> adj(n);
+//!   HLD<0> hld(adj);
+//!   hld.path(u, v, [&](int l, int r) { // [l, r)
+//!   });
+//!   auto [l, r] = hld.subtree(v); // [l, r)
+//! @endcode
+//! @time O(n + q log^2 n)
+//! @space O(n)
+// NOLINTNEXTLINE(readability-identifier-naming)
+template<bool VALS_EDGES> struct HLD {
+  int n;
+  vi p, siz, rt, tin;
+  HLD(auto adj):
+    n(sz(adj)), p(n), siz(n, 1), rt(n), tin(n) {
+    auto dfs1 = [&](auto&& self, int v) -> void {
+      for (int& u : adj[v]) {
+        adj[u].erase(find(all(adj[u]), v));
+        p[u] = v;
+        self(self, u);
+        siz[v] += siz[u];
+        if (siz[u] > siz[adj[v][0]]) swap(u, adj[v][0]);
+      }
+    };
+    dfs1(dfs1, 0);
+    int tim = 0;
+    auto dfs2 = [&](auto&& self, int v) -> void {
+      tin[v] = tim++;
+      for (int u : adj[v]) {
+        rt[u] = (u == adj[v][0] ? rt[v] : u);
+        self(self, u);
+      }
+    };
+    dfs2(dfs2, 0);
+  }
+  void path(int u, int v, auto f) {
+    for (;; v = p[rt[v]]) {
+      if (tin[u] > tin[v]) swap(u, v);
+      if (rt[u] == rt[v]) break;
+      f(tin[rt[v]], tin[v] + 1);
+    }
+    f(tin[u] + VALS_EDGES, tin[v] + 1);
+  }
+  pii subtree(int v) {
+    return {tin[v] + VALS_EDGES, tin[v] + siz[v]};
+  }
+};

tests/library_checker_aizu_tests/data_structures/dsu.test.cpp

@@ -5,12 +5,12 @@ int main() {
   cin.tie(0)->sync_with_stdio(0);
   int n, q;
   cin >> n >> q;
-  dsu uf(n);
+  DSU dsu(n);
   while (q--) {
     int type, u, v;
     cin >> type >> u >> v;
-    if (type == 0) uf.join(u, v);
-    else cout << (uf.go(u) == uf.go(v)) << '\n';
+    if (type == 0) dsu.join(u, v);
+    else cout << (dsu.go(u) == dsu.go(v)) << '\n';
   }
   return 0;
 }

tests/library_checker_aizu_tests/graphs/bcc_callback_lib_checker_two_cc.test.cpp

@@ -14,7 +14,7 @@ int main() {
     adj[u].push_back(v);
     adj[v].push_back(u);
   }
-  dsu uf(n);
+  DSU dsu(n);
   vector<bool> seen(n);
   bcc_callback(adj, [&](const vi& nodes) {
     int count_edges = 0;
@@ -28,10 +28,10 @@ int main() {
       // nodes[0] <=> nodes[1] is a bridge
       return;
     }
-    for (int v : nodes) uf.join(v, nodes[0]);
+    for (int v : nodes) dsu.join(v, nodes[0]);
   });
   vector<vi> two_edge_ccs(n);
-  rep(i, 0, n) two_edge_ccs[uf.go(i)].push_back(i);
+  rep(i, 0, n) two_edge_ccs[dsu.go(i)].push_back(i);
   int cnt_ccs = 0;
   rep(i, 0, n) cnt_ccs += (!empty(two_edge_ccs[i]));
   cout << cnt_ccs << '\n';

tests/library_checker_aizu_tests/graphs/offline_incremental_scc.test.cpp

@@ -20,20 +20,20 @@ int main() {
   iota(all(order), 0);
   ranges::sort(all(order), {},
     [&](int i) { return joins[i]; });
-  dsu uf(n);
+  DSU dsu(n);
   mint sum = 0;
   for (int t = 0, it = 0; t < m; t++) {
     while (it < m && joins[order[it]] <= t) {
       auto [u, v] = eds[order[it]];
-      u = uf.go(u);
-      v = uf.go(v);
-      if (uf.e[u] > uf.e[v]) swap(u, v);
+      u = dsu.go(u);
+      v = dsu.go(v);
+      if (dsu.e[u] > dsu.e[v]) swap(u, v);
       if (u != v) {
         sum = sum + xs[u] * xs[v];
         xs[u] = xs[u] + xs[v];
         xs[v] = xs[u];
       }
-      uf.join(u, v);
+      dsu.join(u, v);
       it++;
     }
     cout << sum.x << '\n';

tests/library_checker_aizu_tests/handmade_tests/dsu.test.cpp

@@ -7,11 +7,11 @@ int main() {
   cin.tie(0)->sync_with_stdio(0);
   for (int n = 0; n <= 100; n++) {
     vector<vector<int>> adj(n);
-    dsu uf(n);
+    DSU dsu(n);
     for (int q = 0; q < n; q++) {
       int u = rnd(0, n - 1);
       int v = rnd(0, n - 1);
-      bool joined = uf.join(u, v);
+      bool joined = dsu.join(u, v);
       {
         vector<bool> vis(n);
         auto dfs = [&](auto&& self, int v) -> void {
@@ -42,7 +42,7 @@ int main() {
         };
         vis[i] = 1;
         dfs(dfs, i);
-        assert(cnt_nodes == uf.size(i));
+        assert(cnt_nodes == dsu.size(i));
       }
     }
   }

tests/library_checker_aizu_tests/trees/hld_aizu1.test.cpp

@@ -0,0 +1,44 @@
+#define PROBLEM \
+  "https://onlinejudge.u-aizu.ac.jp/problems/GRL_5_D"
+#include "../template.hpp"
+#include "../../../library/trees/hld.hpp"
+#include "../../../library/data_structures/bit.hpp"
+int main() {
+  cin.tie(0)->sync_with_stdio(0);
+  int n;
+  cin >> n;
+  vector<vector<int>> adj(n);
+  for (int i = 0; i < n; i++) {
+    int k;
+    cin >> k;
+    for (int j = 0; j < k; j++) {
+      int v;
+      cin >> v;
+      assert(v != 0);
+      adj[i].push_back(v);
+      adj[v].push_back(i);
+    }
+  }
+  HLD<1> hld(adj);
+  BIT bit(n);
+  int q;
+  cin >> q;
+  while (q--) {
+    int type;
+    cin >> type;
+    if (type == 0) {
+      int v, x;
+      cin >> v >> x;
+      assert(v != 0);
+      bit.update(hld.tin[v], x);
+    } else {
+      assert(type == 1);
+      int v;
+      cin >> v;
+      ll sum = 0;
+      hld.path(0, v,
+        [&](int l, int r) { sum += bit.query(l, r); });
+      cout << sum << '\n';
+    }
+  }
+}

tests/library_checker_aizu_tests/trees/hld_aizu2.test.cpp

@@ -0,0 +1,45 @@
+#define PROBLEM \
+  "https://onlinejudge.u-aizu.ac.jp/problems/GRL_5_E"
+#include "../template.hpp"
+#include "../../../library/trees/hld.hpp"
+#include "../../../library/data_structures/lazy_seg_tree.hpp"
+int main() {
+  cin.tie(0)->sync_with_stdio(0);
+  int n;
+  cin >> n;
+  vector<vector<int>> adj(n);
+  for (int i = 0; i < n; i++) {
+    int k;
+    cin >> k;
+    for (int j = 0; j < k; j++) {
+      int v;
+      cin >> v;
+      assert(v != 0);
+      adj[i].push_back(v);
+      adj[v].push_back(i);
+    }
+  }
+  HLD<1> hld(adj);
+  seg_tree st(n);
+  int q;
+  cin >> q;
+  while (q--) {
+    int type;
+    cin >> type;
+    if (type == 0) {
+      int v, x;
+      cin >> v >> x;
+      assert(v != 0);
+      hld.path(0, v,
+        [&](int l, int r) { st.update(l, r, x); });
+    } else {
+      assert(type == 1);
+      int v;
+      cin >> v;
+      ll sum = 0;
+      hld.path(0, v,
+        [&](int l, int r) { sum += st.query(l, r); });
+      cout << sum << '\n';
+    }
+  }
+}

tests/library_checker_aizu_tests/trees/hld_lib_checker_path.test.cpp

@@ -0,0 +1,38 @@
+#define PROBLEM \
+  "https://judge.yosupo.jp/problem/vertex_add_path_sum"
+#include "../template.hpp"
+#include "../../../library/trees/hld.hpp"
+#include "../../../library/data_structures/bit.hpp"
+int main() {
+  cin.tie(0)->sync_with_stdio(0);
+  int n, q;
+  cin >> n >> q;
+  vector<int> a(n);
+  for (int i = 0; i < n; i++) cin >> a[i];
+  vector<vector<int>> adj(n);
+  for (int i = 0; i < n - 1; i++) {
+    int u, v;
+    cin >> u >> v;
+    adj[u].push_back(v);
+    adj[v].push_back(u);
+  }
+  HLD<0> hld(adj);
+  BIT bit(n);
+  for (int i = 0; i < n; i++) bit.update(hld.tin[i], a[i]);
+  while (q--) {
+    int type;
+    cin >> type;
+    if (type == 0) {
+      int v, x;
+      cin >> v >> x;
+      bit.update(hld.tin[v], x);
+    } else {
+      int u, v;
+      cin >> u >> v;
+      ll sum = 0;
+      hld.path(u, v,
+        [&](int l, int r) { sum += bit.query(l, r); });
+      cout << sum << '\n';
+    }
+  }
+}

tests/library_checker_aizu_tests/trees/hld_lib_checker_subtree.test.cpp

@@ -0,0 +1,42 @@
+#define PROBLEM \
+  "https://judge.yosupo.jp/problem/vertex_add_subtree_sum"
+#include "../template.hpp"
+#include "../../../library/trees/hld.hpp"
+#include "../../../library/data_structures/bit.hpp"
+int main() {
+  cin.tie(0)->sync_with_stdio(0);
+  int n, q;
+  cin >> n >> q;
+  vector<ll> a(n);
+  for (int i = 0; i < n; i++) cin >> a[i];
+  vector<vector<int>> adj(n);
+  for (int i = 1; i < n; i++) {
+    int par;
+    cin >> par;
+    adj[par].push_back(i);
+    adj[i].push_back(par);
+  }
+  HLD<0> hld_nodes(adj);
+  HLD<1> hld_edges(adj);
+  BIT bit(n);
+  for (int i = 0; i < n; i++)
+    bit.update(hld_nodes.tin[i], a[i]);
+  while (q--) {
+    int type;
+    cin >> type;
+    if (type == 0) {
+      int v, x;
+      cin >> v >> x;
+      bit.update(hld_nodes.tin[v], x);
+      a[v] += x;
+    } else {
+      int v;
+      cin >> v;
+      auto [l, r] = hld_nodes.subtree(v);
+      ll res = bit.query(l, r);
+      auto [l_edges, r_edges] = hld_edges.subtree(v);
+      assert(res == bit.query(l_edges, r_edges) + a[v]);
+      cout << res << '\n';
+    }
+  }
+}