lca with precomp dfs order (#34)

Author: cameroncuster

examples/graphs/lca_yosupo.rs

@@ -9,36 +9,35 @@ fn main() {
         q: usize,
     }
 
-    // forest with an undirected tree and a rooted directed tree
-    let mut undir_for = vec![vec![]; n + n];
+    // undirected tree and a rooted directed tree
+    let mut undir = vec![vec![]; n + n];
     let mut dir = vec![vec![]; n];
     for c in 1..n {
         input! {
             p: usize,
         }
 
         // undirected tree
-        for &(u, v) in [(p, c), (p + n, c + n)].iter() {
-            undir_for[u].push(v);
-            undir_for[v].push(u);
-        }
+        undir[p].push(c);
+        undir[c].push(p);
 
         // directed tree
         dir[p].push(c);
     }
 
-    let undir_lca = LCA::new(&undir_for);
+    let undir_lca = LCA::new(&undir);
     let dir_lca = LCA::new(&dir);
     for _ in 0..q {
         input! {
             u: usize,
             v: usize,
         }
-        let res_undir1 = undir_lca.lca(u, v);
-        let res_undir2 = undir_lca.lca(u + n, v + n) - n;
+
+        let res_undir = undir_lca.lca(u, v);
         let res_dir = dir_lca.lca(u, v);
-        assert_eq!(res_undir1, res_undir2);
-        assert_eq!(res_undir1, res_dir);
-        println!("{}", res_dir);
+
+        assert_eq!(res_undir, res_dir);
+
+        println!("{}", res_undir);
     }
 }

src/graphs/dfs_order.rs

@@ -0,0 +1,14 @@
+pub fn dfs_order(adj: &[Vec<usize>]) -> Vec<usize> {
+    let n = adj.len();
+    let mut order = Vec::with_capacity(n);
+    fn dfs(u: usize, p: usize, adj: &[Vec<usize>], order: &mut Vec<usize>) {
+        order.push(u);
+        for &v in &adj[u] {
+            if v != p {
+                dfs(v, u, adj, order);
+            }
+        }
+    }
+    dfs(0, 0, adj, &mut order);
+    order
+}

src/graphs/lca.rs

@@ -1,6 +1,7 @@
 //! # Lowest Common Ancestor
 
 use crate::data_structures::rmq::RMQ;
+use crate::graphs::dfs_order::dfs_order;
 
 /// # Example
 /// ```
@@ -23,48 +24,33 @@ use crate::data_structures::rmq::RMQ;
 /// ```
 pub struct LCA {
     tin: Vec<usize>,
-    p: Vec<Option<usize>>,
+    p: Vec<usize>,
     rmq: RMQ<(usize, usize)>,
 }
 
 impl LCA {
     /// Create a new LCA struct
     ///
-    /// `adj` can be undirected tree, directed tree (rooted at node 0), or a forest of undirected
-    /// trees
+    /// `adj` can be undirected tree or a directed tree (rooted at node 0)
     ///
     /// # Complexity (n = adj.len())
     /// - Time: O(n log n)
     /// - Space: O(n log n)
     pub fn new(adj: &[Vec<usize>]) -> Self {
         let n = adj.len();
-        let mut d = vec![0; n];
         let mut tin = vec![0; n];
-        let mut p = vec![None; n];
-        let mut order = Vec::with_capacity(n);
-        fn dfs(
-            u: usize,
-            p: &mut [Option<usize>],
-            adj: &[Vec<usize>],
-            d: &mut [usize],
-            tin: &mut [usize],
-            order: &mut Vec<usize>,
-        ) {
-            tin[u] = order.len();
-            order.push(u);
+        let mut p = vec![0; n];
+        let mut d = vec![0; n];
+        let order = dfs_order(adj);
+        for (i, &u) in order.iter().enumerate() {
+            tin[u] = i;
             for &v in &adj[u] {
-                if p[u] != Some(v) {
+                if v != p[u] {
+                    p[v] = u;
                     d[v] = d[u] + 1;
-                    p[v] = Some(u);
-                    dfs(v, p, adj, d, tin, order);
                 }
             }
         }
-        for s in 0..n {
-            if p[s].is_none() {
-                dfs(s, &mut p, adj, &mut d, &mut tin, &mut order);
-            }
-        }
         let d_with_order: Vec<(usize, usize)> = order.iter().map(|&u| (d[u], u)).collect();
         let rmq = RMQ::new(&d_with_order, std::cmp::min);
         LCA { tin, p, rmq }
@@ -83,6 +69,6 @@ impl LCA {
         if le > ri {
             std::mem::swap(&mut le, &mut ri);
         }
-        self.p[self.rmq.query(le + 1..ri + 1).1].unwrap()
+        self.p[self.rmq.query(le + 1..ri + 1).1]
     }
 }

src/graphs/mod.rs

@@ -1,3 +1,4 @@
 //! # Graph Algorithms
+mod dfs_order;
 pub mod dijk;
 pub mod lca;