adding new kth par alg

Author: lrvideckis

library/trees/uncommon/ladder_decomposition.hpp

@@ -30,8 +30,8 @@ struct ladder {
           if (sz(x) > sz(path)) swap(x, path);
           for (int y : x) idx[y] = i;
           for (int y : x) lad[i++] = y;
-          rep(j, 0, min<int>(sz(x), d[v]))
-            lad[i++] = st[d[u] - j];
+          rep(k, 0, min<int>(sz(x), d[v]))
+            lad[i++] = st[d[u] - k];
         }
       path.push_back(u);
       return path;

library/trees/uncommon/linear_kth_par.hpp

@@ -0,0 +1,56 @@
+#pragma once
+//! https://link.springer.com/chapter/10.1007/978-3-032-29003-8_29
+//! @code
+//!   vector<basic_string<int>> g(n);
+//!   linear_kth_par kp(g);
+//!   kp.kth_par(v, k); // k edges up from v
+//!   kp.kth_par(v, 1); // v's parent
+//! @endcode
+//! @time O(n + q)
+//! @space O(n)
+struct linear_kth_par {
+  int n;
+  vi d, et_i, et_d, idx, lad;
+  vector<basic_string<int>> jmp;
+  linear_kth_par(const auto& g):
+    n(sz(g)), d(n), et_i(n), et_d(2 * n), idx(n),
+    lad(2 * n), jmp(2 * n) {
+    int i = 0, j = 1;
+    vi st(n);
+    auto calc = [&](int u) {
+      et_d[et_i[u] = j] = d[u];
+      for (int k = 1; k <= min(j & -j, d[u]); k *= 2)
+        jmp[j] += st[d[u] - k];
+      j++;
+    };
+    auto dfs = [&](auto dfs, int u, int p) -> vi {
+      vi path;
+      st[d[u]] = u;
+      calc(u);
+      for (int v : g[u])
+        if (v != p) {
+          d[v] = d[u] + 1;
+          vi x = dfs(dfs, v, u);
+          calc(u);
+          if (sz(x) > sz(path)) swap(x, path);
+          for (int y : x) idx[y] = i;
+          for (int y : x) lad[i++] = y;
+          rep(k, 0, min<int>(sz(x), d[v]))
+            lad[i++] = st[d[u] - k];
+        }
+      path.push_back(u);
+      return path;
+    };
+    vi x = dfs(dfs, 0, 0);
+    for (int y : x) idx[y] = i;
+    for (int y : x) lad[i++] = y;
+  }
+  int kth_par(int u, int k) {
+    assert(0 <= k && k <= d[u]);
+    if (k == 0) return u;
+    int anc_d = d[u] - k, bc = bit_ceil(k + 0u);
+    int j = (et_i[u] + bc / 2) & -bc;
+    int i = idx[jmp[j].at(__lg(et_d[j] - anc_d))];
+    return lad[i + d[lad[i]] - anc_d];
+  }
+};

tests/library_checker_aizu_tests/handmade_tests/kth_par.test.cpp

@@ -0,0 +1,34 @@
+#define PROBLEM \
+  "https://onlinejudge.u-aizu.ac.jp/problems/ITP1_1_A"
+#include "../template.hpp"
+#include "../../../library/contest/random.hpp"
+#include "../../../library/trees/uncommon/hagerup_kth_par.hpp"
+#include "../../../library/trees/uncommon/linear_kth_par.hpp"
+int main() {
+  for (int n = 1; n <= 600; n++) {
+    for (int param = 1; param < 20; param++) {
+      vector<int> par(n);
+      vector<vector<int>> adj(n);
+      for (int i = 1; i < n; i++) {
+        if (param == 19) par[i] = rnd(0, i - 1);
+        else par[i] = rnd(max(0, i - param), i - 1);
+        adj[par[i]].push_back(i);
+        adj[i].push_back(par[i]);
+      }
+      hagerup<2> hagerup_kth_par(adj);
+      linear_kth_par lin_kth_par(adj);
+      for (int i = 0; i < n; i++) {
+        for (int k = 0, kth_parent_naive = i;
+          k <= hagerup_kth_par.d[i];
+          k++, kth_parent_naive = par[kth_parent_naive]) {
+          assert(kth_parent_naive ==
+                 hagerup_kth_par.kth_par(i, k));
+          assert(
+            kth_parent_naive == lin_kth_par.kth_par(i, k));
+        }
+      }
+    }
+  }
+  cout << "Hello World\n";
+  return 0;
+}

tests/library_checker_aizu_tests/trees/kth_path_hagerup.test.cpp

@@ -0,0 +1,44 @@
+#define PROBLEM \
+  "https://judge.yosupo.jp/problem/jump_on_tree"
+#include "../template.hpp"
+#include "../../../library/trees/linear_lca.hpp"
+#include "../../../library/trees/uncommon/hagerup_kth_par.hpp"
+#include "../../../library/trees/lca_rmq.hpp"
+#include "../compress_tree_asserts.hpp"
+int main() {
+  cin.tie(0)->sync_with_stdio(0);
+  int n, q;
+  cin >> n >> q;
+  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);
+  }
+  linear_lca lin_lca(adj);
+  hagerup<1> lin_kp_1(adj);
+  hagerup<2> lin_kp_2(adj);
+  hagerup<3> lin_kp_3(adj);
+  hagerup<4> lin_kp_4(adj);
+  LCA lc(adj);
+  compress_tree_asserts(adj, lc);
+  auto get_kth_par = [&](int v, int k) -> int {
+    int res = lin_kp_1.kth_par(v, k);
+    assert(res == lin_kp_2.kth_par(v, k));
+    assert(res == lin_kp_3.kth_par(v, k));
+    assert(res == lin_kp_4.kth_par(v, k));
+    return res;
+  };
+  while (q--) {
+    int u, v, k;
+    cin >> u >> v >> k;
+    int lca_d = lin_kp_2.d[lin_lca.lca(u, v)];
+    int u_lca = lin_kp_2.d[u] - lca_d;
+    int v_lca = lin_kp_2.d[v] - lca_d;
+    if (k <= u_lca) cout << get_kth_par(u, k) << '\n';
+    else if (k <= u_lca + v_lca)
+      cout << get_kth_par(v, u_lca + v_lca - k) << '\n';
+    else cout << -1 << '\n';
+  }
+}

tests/library_checker_aizu_tests/trees/kth_path_linear.test.cpp

@@ -2,9 +2,7 @@
   "https://judge.yosupo.jp/problem/jump_on_tree"
 #include "../template.hpp"
 #include "../../../library/trees/linear_lca.hpp"
-#include "../../../library/trees/uncommon/hagerup_kth_par.hpp"
-#include "../../../library/trees/lca_rmq.hpp"
-#include "../compress_tree_asserts.hpp"
+#include "../../../library/trees/uncommon/linear_kth_par.hpp"
 int main() {
   cin.tie(0)->sync_with_stdio(0);
   int n, q;
@@ -17,28 +15,18 @@ int main() {
     adj[v].push_back(u);
   }
   linear_lca lin_lca(adj);
-  hagerup<1> lin_kp_1(adj);
-  hagerup<2> lin_kp_2(adj);
-  hagerup<3> lin_kp_3(adj);
-  hagerup<4> lin_kp_4(adj);
-  LCA lc(adj);
-  compress_tree_asserts(adj, lc);
-  auto get_kth_par = [&](int v, int k) -> int {
-    int res = lin_kp_1.kth_par(v, k);
-    assert(res == lin_kp_2.kth_par(v, k));
-    assert(res == lin_kp_3.kth_par(v, k));
-    assert(res == lin_kp_4.kth_par(v, k));
-    return res;
-  };
+  linear_kth_par lin_kth_par(adj);
   while (q--) {
     int u, v, k;
     cin >> u >> v >> k;
-    int lca_d = lin_kp_2.d[lin_lca.lca(u, v)];
-    int u_lca = lin_kp_2.d[u] - lca_d;
-    int v_lca = lin_kp_2.d[v] - lca_d;
-    if (k <= u_lca) cout << get_kth_par(u, k) << '\n';
+    int lca_d = lin_lca.d[lin_lca.lca(u, v)];
+    int u_lca = lin_lca.d[u] - lca_d;
+    int v_lca = lin_lca.d[v] - lca_d;
+    if (k <= u_lca)
+      cout << lin_kth_par.kth_par(u, k) << '\n';
     else if (k <= u_lca + v_lca)
-      cout << get_kth_par(v, u_lca + v_lca - k) << '\n';
+      cout << lin_kth_par.kth_par(v, u_lca + v_lca - k)
+           << '\n';
     else cout << -1 << '\n';
   }
 }