update ladder decomp

Author: lrvideckis

library/trees/uncommon/ladder_decomposition.hpp

@@ -4,53 +4,53 @@
 //! https://youtu.be/0rCFkuQS968
 //! @code
 //!   ladder ld(g);
+//!   ld.kth_par(u, k); // kth parent of u
+//!   ld.kth_par(u, 0); // u
+//!   ld.kth_par(u, 1); // p[u]
 //!   // KACTL functions
 //!   int kth_par = jmp(ld.jmp, u, k);
 //!   int curr_lca = lca(ld.jmp, ld.d, u, v);
 //! @endcode
 struct ladder {
   int n;
-  vi d, p, leaf, idx, lad;
+  vi d, p, idx, l;
   vector<vi> jmp;
   //! @param g forest (rooted or unrooted)
   //! @time O(n log n)
   //! @space O(n log n) for jmp. Everything else is O(n)
   ladder(const auto& g):
-    n(sz(g)), d(n), p(n), leaf(n), idx(n), lad(2 * n) {
-    auto dfs = [&](auto dfs, int u) -> void {
-      leaf[u] = u;
+    n(sz(g)), d(n), p(n), idx(n), l(2 * n) {
+    int i = 0;
+    vi s(n);
+    auto dfs = [&](auto dfs, int u) -> vi {
+      vi path;
+      s[d[u]] = u;
       for (int v : g[u])
         if (v != p[u]) {
           d[v] = d[p[v] = u] + 1;
-          dfs(dfs, v);
-          if (d[leaf[u]] < d[leaf[v]]) leaf[u] = leaf[v];
+          vi x = dfs(dfs, v);
+          if (sz(x) > sz(path)) swap(x, path);
+          for (int y : x) idx[y] = i;
+          for (int y : x) l[i++] = y;
+          rep(j, 0, min(sz(x), d[v])) l[i++] = s[d[u] - j];
         }
+      path.push_back(u);
+      return path;
     };
-    dfs(dfs, 0);
-    int pos = 0;
-    rep(i, 0, n) if (p[i] == i || leaf[p[i]] != leaf[i]) {
-      int l = leaf[i];
-      int len = min((d[l] - d[i]) * 2, d[l] + 1);
-      idx[l] = pos;
-      for (; len--; l = p[l]) lad[pos++] = l;
-    }
+    vi x = dfs(dfs, 0);
+    for (int y : x) idx[y] = i;
+    for (int y : x) l[i++] = y;
     jmp = treeJump(p);
   }
-  //! @param u query node
-  //! @param k number of edges
-  //! @returns a node k edges up from u. With k=1, this
-  //! returns u's parent.
-  //! @time O(1)
-  //! @space O(1)
   int kth_par(int u, int k) {
     assert(0 <= k && k <= d[u]);
     if (k == 0) return u;
     int bit = __lg(k);
     u = jmp[bit][u], k -= (1 << bit);
-    int l = idx[leaf[u]] + d[leaf[u]] - d[u];
-    assert(lad[l] == u);
-    // subarray [l, l+k] of lad corresponds to the rest
+    int i = idx[u], j = i + d[l[i]] - d[u];
+    assert(l[j] == u);
+    // subarray [j, j+k] of l corresponds to the rest
     // of the jump
-    return lad[l + k];
+    return l[j + k];
   }
 };

tests/library_checker_aizu_tests/trees/kth_path_ladder.test.cpp

@@ -1,8 +1,8 @@
 #define PROBLEM \
   "https://judge.yosupo.jp/problem/jump_on_tree"
 #include "../template.hpp"
+#define sz(x) (int)ssize(x)
 #include "../../../library/trees/uncommon/ladder_decomposition.hpp"
-#include "../../../library/trees/uncommon/linear_kth_par.hpp"
 int main() {
   cin.tie(0)->sync_with_stdio(0);
   int n, q;