Next on path (#96)

* convert

* converted RMQ

* fix

* fix

* nit

* actually let's do fnmut

* Revert "actually let's do fnmut"

This reverts commit 946ebfd.

* saving progress

* fix

* fix docs

* add extra examples

* nit

* fix bug

* actually this is better

* Update lca.rs

---------

Co-authored-by: Luke Videckis <lukevideckis@gmail.com>
Co-authored-by: Cameron Custer <73217097+cameroncuster@users.noreply.github.com>

Author: lrvideckis

Cargo.toml

@@ -215,3 +215,7 @@ path = "examples/monotonic/count_rects.rs"
 [[example]]
 name = "cartesian_tree"
 path = "examples/monotonic/cartesian_tree.rs"
+
+[[example]]
+name = "lca_rmq_next_on_path"
+path = "examples/graphs/lca_rmq_next_on_path.rs"

examples/graphs/lca_rmq_next_on_path.rs

@@ -0,0 +1,65 @@
+// verification-helper: PROBLEM https://judge.yosupo.jp/problem/tree_diameter
+
+use proconio::input;
+use programming_team_code_rust::graphs::lca::LCA;
+
+fn main() {
+    input! {
+        n: usize,
+    }
+
+    let mut adj = vec![vec![]; n];
+    let mut adj_w = vec![vec![]; n];
+    for _ in 1..n {
+        input! {
+            u: usize,
+            v: usize,
+            w: i64,
+        }
+        adj[u].push(v);
+        adj[v].push(u);
+        adj_w[u].push((v, w));
+        adj_w[v].push((u, w));
+    }
+
+    let lca = LCA::new(&adj);
+
+    let mut dist_root = vec![0; n];
+    fn dfs(u: usize, p: Option<usize>, adj_w: &[Vec<(usize, i64)>], dist_root: &mut Vec<i64>) {
+        for &(v, w) in &adj_w[u] {
+            if Some(v) == p {
+                continue;
+            }
+            dist_root[v] = w + dist_root[u];
+            dfs(v, Some(u), adj_w, dist_root);
+        }
+    }
+
+    dfs(0, None, &adj_w, &mut dist_root);
+
+    let dist_path =
+        |u: usize, v: usize| -> i64 { dist_root[u] + dist_root[v] - 2 * dist_root[lca.lca(u, v)] };
+
+    let mut u = 0;
+    for i in 1..n {
+        if dist_path(0, u) < dist_path(0, i) {
+            u = i;
+        }
+    }
+
+    let mut v = 0;
+    for i in 1..n {
+        if dist_path(u, v) < dist_path(u, i) {
+            v = i;
+        }
+    }
+
+    println!("{} {}", dist_path(u, v), lca.dist(u, v) + 1);
+
+    print!("{} ", u);
+    while u != v {
+        u = lca.next_on_path(u, v);
+        print!(" {}", u);
+    }
+    println!();
+}

src/graphs/lca.rs

@@ -3,7 +3,7 @@
 use crate::data_structures::rmq::RMQ;
 use crate::graphs::dfs_order::get_dfs_preorder;
 
-type OpType = fn(&(usize, usize), &(usize, usize)) -> (usize, usize);
+type Pair = (usize, usize);
 
 /// # Example
 /// ```
@@ -19,11 +19,17 @@ type OpType = fn(&(usize, usize), &(usize, usize)) -> (usize, usize);
 /// let lca = LCA::new(&adj);
 /// assert_eq!(lca.lca(1, 3), 0);
 /// assert_eq!(lca.lca(2, 3), 2);
+/// assert_eq!(lca.dist(1, 3), 3);
+/// assert_eq!(lca.in_sub(2, 3), true);
+/// assert_eq!(lca.next_on_path(0, 3), 2);
+/// assert!(std::panic::catch_unwind(|| lca.next_on_path(1, 1)).is_err());
 /// ```
 pub struct LCA {
     tin: Vec<usize>,
     p: Vec<Option<usize>>,
-    rmq: RMQ<(usize, usize), OpType>,
+    d: Vec<usize>,
+    siz: Vec<usize>,
+    rmq: RMQ<Pair, fn(&Pair, &Pair) -> Pair>,
 }
 
 impl LCA {
@@ -49,13 +55,21 @@ impl LCA {
                 }
             }
         }
+        let mut siz = vec![1; n];
+        for &u in order.iter().rev() {
+            if let Some(par) = p[u] {
+                siz[par] += siz[u];
+            }
+        }
         LCA {
-            tin,
             p,
             rmq: RMQ::new(
                 &order.iter().map(|&u| (d[u], u)).collect::<Vec<_>>(),
-                |&x, &y| std::cmp::min(x, y),
+                |&x, &y| if x.0 < y.0 { x } else { y },
             ),
+            d,
+            tin,
+            siz,
         }
     }
 
@@ -74,4 +88,36 @@ impl LCA {
         }
         self.p[self.rmq.query(le + 1..ri + 1).1].unwrap()
     }
+
+    /// Gets number of edges on path from u to v
+    ///
+    /// # Complexity
+    /// - Time: O(1)
+    /// - Space: O(1)
+    pub fn dist(&self, u: usize, v: usize) -> usize {
+        self.d[u] + self.d[v] - 2 * self.d[self.lca(u, v)]
+    }
+
+    /// Returns true iff v is in u's subtree
+    ///
+    /// # Complexity
+    /// - Time: O(1)
+    /// - Space: O(1)
+    pub fn in_sub(&self, u: usize, v: usize) -> bool {
+        (self.tin[u]..self.tin[u] + self.siz[u]).contains(&self.tin[v])
+    }
+
+    /// Gets \[u, p\[u\], .., lca(u,v), .., p\[v\], v\]\[1\]
+    ///
+    /// # Complexity
+    /// - Time: O(1)
+    /// - Space: O(1)
+    pub fn next_on_path(&self, u: usize, v: usize) -> usize {
+        assert!(u != v);
+        if self.in_sub(u, v) {
+            self.rmq.query(self.tin[u] + 1..self.tin[v] + 1).1
+        } else {
+            self.p[u].unwrap()
+        }
+    }
 }