Linear kth par nit (#150)

* switch test

* [auto-verifier] verify commit c139baa

* add changes

* [auto-verifier] verify commit d3d4c46

* speed up test

* [auto-verifier] verify commit a7b7c0f

* convert back to vector

* [auto-verifier] verify commit ffd205e

* [auto-verifier] verify commit 738bbf1

* final golf

* adding git ignore to try to fix CI

* final nits

* [auto-verifier] remove .verify-helper/.gitignore (see online-judge-tools/verification-helper#332)

* nit

---------

Co-authored-by: GitHub <[email protected]>

Author: lrvideckis

library/trees/ladder_decomposition/linear_kth_par.hpp

@@ -7,62 +7,47 @@
 //!   }
 //!   vector<basic_string<int>> adj(n);
 //!   linear_kth_par kp(adj);
-//!   int kth_par = kp.kth_par(v, k);
+//!   kp.kth_par(v, k); // k edges up from v
+//!   kp.kth_par(v, 1); // v's parent
 //! @endcode
-//! kth_par = a node k edges up from v
-//! @time O(n + q)
-//! @space O(n)
-struct linear_kth_par {
-  struct node {
-    int d, p = -1, dl, idx_j;
-    basic_string<int> lad;
-  };
-  vector<node> t;
-  vector<pii> j;
+//! @time O(n*max((2*KAPPA+3)/KAPPA,2*KAPPA) + q)
+//! @space O(n*max((2*KAPPA+3)/KAPPA,2*KAPPA))
+template<int KAPPA = 2> struct linear_kth_par {
+  int n;
+  vi d, leaf, pos, jmp;
+  vector<vi> lad;
   linear_kth_par(const auto& adj):
-    t(sz(adj)), j(2 * sz(t)) {
-    vi st;
-    int pos = 1;
-    auto add_j = [&]() -> void {
-      j[pos] = {
-        st[max<int>(0, sz(st) - 1 - 2 * (pos & -pos))],
-        st[max<int>(0, sz(st) - 1 - 4 * (pos & -pos))]};
-      pos++;
+    n(sz(adj)), d(n), leaf(n), pos(n), jmp(2 * n), lad(n) {
+    static_assert(KAPPA >= 1);
+    int t = 1;
+    vi st(n);
+    auto calc = [&](int siz) {
+      jmp[t] = st[max(0, siz - KAPPA * (t & -t))];
+      t++;
     };
-    auto dfs = [&](auto&& self, int v) -> void {
-      st.push_back(v);
-      t[v].idx_j = pos, t[v].dl = v;
-      add_j();
+    auto dfs = [&](auto&& self, int v, int p) {
+      st[d[v]] = v;
+      int& l = leaf[v] = v;
+      pos[v] = t;
+      calc(d[v]);
       for (int u : adj[v])
-        if (u != t[v].p) {
-          t[u].d = t[t[u].p = v].d + 1;
-          self(self, u);
-          if (t[t[u].dl].d > t[t[v].dl].d)
-            t[v].dl = t[u].dl;
-          add_j();
+        if (u != p) {
+          d[u] = 1 + d[v];
+          self(self, u, v);
+          if (d[l] < d[leaf[u]]) l = leaf[u];
+          calc(d[v]);
         }
-      st.pop_back();
+      int s = (d[l] - d[v]) * (2 * KAPPA + 3) / KAPPA;
+      s = min(max(s, 2 * KAPPA), d[l] + 1);
+      rep(i, sz(lad[l]), s) lad[l].push_back(st[d[l] - i]);
     };
-    rep(i, 0, sz(t)) {
-      if (t[i].p == -1) dfs(dfs, i);
-      if (t[i].p == -1 || t[t[i].p].dl != t[i].dl) {
-        int v = t[i].dl, len = (t[v].d - t[i].d) * 2;
-        auto& lad = t[v].lad;
-        for (; v != -1 && len--; v = t[v].p) lad += v;
-      }
-    }
+    dfs(dfs, 0, 0);
   }
   int kth_par(int v, int k) {
-    assert(0 <= k && k <= t[v].d);
-    switch (k) {
-    case 0: return v;
-    case 1: return t[v].p;
-    case 2: return t[t[v].p].p;
-    default:
-      int i = bit_floor(unsigned(k / 3));
-      auto [j1, j2] = j[(t[v].idx_j & -i) | i];
-      int leaf = t[t[v].d - t[j2].d <= k ? j2 : j1].dl;
-      return t[leaf].lad[k + t[leaf].d - t[v].d];
-    }
+    assert(0 <= k && k <= d[v]);
+    int j = v;
+    if (unsigned b = k / (KAPPA + 1); b)
+      b = bit_floor(b), j = jmp[(pos[v] & -b) | b];
+    return j = leaf[j], lad[j][k + d[j] - d[v]];
   }
 };

library/trees/linear_kth_path.hpp

@@ -16,14 +16,6 @@ int main() {
   }
   ladder ld(adj);
   assert(sz(ld.l_tbl) <= 2 * n - ld.d[0]);
-  vector<vector<int>> adj_rooted(n + n);
-  for (int i = 0; i < n; i++)
-    if (ld.p[i] != i) {
-      adj_rooted[ld.p[i]].push_back(i);
-      adj_rooted[ld.p[i] + n].push_back(i + n);
-    }
-  ladder ld_rooted(adj_rooted);
-  linear_kth_par lin_ld_rooted(adj_rooted);
   while (q--) {
     int u, v, k;
     cin >> u >> v >> k;
@@ -34,23 +26,10 @@ int main() {
     else if (k <= u_lca) {
       int res = ld.kth_par(u, k);
       assert(res == jmp(ld.b_tbl, u, k));
-      assert(res == ld_rooted.kth_par(u, k));
-      assert(res == ld_rooted.kth_par(u + n, k) - n);
-      assert(res == lin_ld_rooted.kth_par(u, k));
-      assert(res == lin_ld_rooted.kth_par(u + n, k) - n);
       cout << res << '\n';
     } else {
       int res = ld.kth_par(v, u_lca + v_lca - k);
       assert(res == jmp(ld.b_tbl, v, u_lca + v_lca - k));
-      assert(
-        res == ld_rooted.kth_par(v, u_lca + v_lca - k));
-      assert(res ==
-        ld_rooted.kth_par(v + n, u_lca + v_lca - k) - n);
-      assert(res ==
-        lin_ld_rooted.kth_par(v, u_lca + v_lca - k));
-      assert(res ==
-        lin_ld_rooted.kth_par(v + n, u_lca + v_lca - k) -
-          n);
       cout << res << '\n';
     }
   }

tests/library_checker_aizu_tests/trees/kth_path_ladder.test.cpp

@@ -1,7 +1,8 @@
 #define PROBLEM \
   "https://judge.yosupo.jp/problem/jump_on_tree"
 #include "../template.hpp"
-#include "../../../library/trees/linear_kth_path.hpp"
+#include "../../../library/trees/linear_lca.hpp"
+#include "../../../library/trees/ladder_decomposition/linear_kth_par.hpp"
 #include "../../../library/trees/lca_rmq/lca_rmq.hpp"
 #include "../compress_tree_asserts.hpp"
 int main() {
@@ -15,12 +16,29 @@ int main() {
     adj[u].push_back(v);
     adj[v].push_back(u);
   }
-  linear_kth_path lin_kth_path(adj);
+  linear_lca lin_lca(adj);
+  linear_kth_par<1> lin_kp_1(adj);
+  linear_kth_par<2> lin_kp_2(adj);
+  linear_kth_par<3> lin_kp_3(adj);
+  linear_kth_par<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;
-    cout << lin_kth_path.kth_path(u, v, k) << '\n';
+    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';
   }
 }