add changes

Author: lrvideckis

library/trees/ladder_decomposition/linear_kth_par.hpp

@@ -7,62 +7,53 @@
 //!   }
 //!   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*(2*KAPPA+1)/(KAPPA-1) + n*2*(KAPPA-1) + q)
+//! @space O(n*(2*KAPPA+1)/(KAPPA-1) + n*2*(KAPPA-1))
+template<int KAPPA = 3> struct linear_kth_par {
+  int n;
+  vi d, leaf, pos, jmp;
+  vector<basic_string<int>> 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 >= 2);
+    int t = 1;
+    vi st(n);
+    auto calc = [&](int siz) {
+      jmp[t] = st[max(0, siz - (KAPPA - 1) * (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) -> void {
+      st[d[v]] = 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[leaf[v]] < d[leaf[u]]) leaf[v] = leaf[u];
+          calc(d[v]);
         }
-      st.pop_back();
+      // TODO golf this further
+      int len = (d[leaf[v]] - d[v]) * (2 * KAPPA + 1) /
+        (KAPPA - 1);
+      len = max(len, 2 * (KAPPA - 1));
+      int i = d[leaf[v]] - sz(lad[leaf[v]]);
+      while (i && sz(lad[leaf[v]]) < len)
+        lad[leaf[v]] += st[i--];
+      assert(lad[leaf[v]].back() > 0);
     };
-    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; b)
+      b = bit_floor(b), j = jmp[(pos[v] & -b) | b];
+    // TODO array index out of bounds goes uncaught here
+    // because of basic_string
+    return j = leaf[j], lad[j][k + d[j] - d[v]];
   }
 };

tests/library_checker_aizu_tests/trees/kth_path_linear.test.cpp

@@ -17,18 +17,28 @@ int main() {
     adj[v].push_back(u);
   }
   linear_lca lin_lca(adj);
-  linear_kth_par lin_kth_par(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);
+  linear_kth_par<5> lin_kp_5(adj);
   LCA lc(adj);
   compress_tree_asserts(adj, lc);
+  auto get_kth_par = [&](int v, int k) -> int {
+    int 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));
+    assert(res == lin_kp_5.kth_par(v, k));
+    return res;
+  };
   while (q--) {
     int u, v, k;
     cin >> u >> v >> k;
-    int lca_d = lin_kth_par.t[lin_lca.lca(u, v)].d;
-    int u_lca = lin_kth_par.t[u].d - lca_d;
-    int v_lca = lin_kth_par.t[v].d - lca_d;
-    if (k <= u_lca) cout << lin_kth_par.kth_par(u, 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 << lin_kth_par.kth_par(v, u_lca + v_lca - k) << '\n';
+      cout << get_kth_par(v, u_lca + v_lca - k) << '\n';
     else cout << -1 << '\n';
   }
 }