refactor: move LCA with binary lifting to the project

Author: ThiagoFBastos

competitive_programming/src/graphs/binary_lifting.rs

@@ -0,0 +1,123 @@
+#[derive(Clone)]
+pub struct LCA {
+    adj: Vec<Vec<(usize, i64)>>, // the adjacency list with weights
+    depth: Vec<i32>,             // the depth/height of each vertex
+    st: Vec<Vec<usize>>,         // the sparse table for jumps
+    weight: Vec<i64>,            // the sum of weights from the root to each vertex
+}
+
+impl LCA {
+    /**
+     * Create a new instance of LCA data structure
+     * @param n the number of nodes
+     * @return the new LCA data structure
+     */
+    pub fn new(n: usize) -> Self {
+        assert_ne!(n, 0);
+
+        let lg = 32 - (n as i32).leading_zeros();
+
+        Self {
+            adj: vec![vec![]; n],
+            depth: vec![0; n],
+            st: vec![vec![0; lg as usize]; n],
+            weight: vec![0; n],
+        }
+    }
+
+    /**
+     * Do a DFS from the root to the vertexes to calculate some properties
+     * @param u the vertex to do the search
+     * @param p the parent of vertex u
+     */
+    fn dfs(&mut self, u: usize, p: usize) {
+        self.st[u][0] = p;
+
+        for i in 1..self.st[u].len() {
+            self.st[u][i] = self.st[self.st[u][i - 1]][i - 1];
+        }
+
+        for &(v, w) in self.adj[u].clone().iter() {
+            if v == p {
+                continue;
+            }
+
+            self.depth[v] = 1 + self.depth[u];
+            self.weight[v] = w + self.weight[u];
+            self.dfs(v, u);
+        }
+    }
+
+    /**
+     * Build the LCA data structure using a vertex as root
+     * @param root the root of the tree
+     */
+    pub fn build(&mut self, root: usize) {
+        self.depth[root] = 0;
+        self.weight[root] = 0;
+        self.dfs(root, root);
+    }
+
+    /**
+     * Add an undirected edge between u and v with weight w, being that w = 1 if the tree is unweighted
+     * @param u a vertex of the edge
+     * @param v a vertex of the edge
+     * @param w the weight of the edge
+     */
+    pub fn add_edge(&mut self, u: usize, v: usize, w: i64) {
+        self.adj[u].push((v, w));
+        self.adj[v].push((u, w));
+    }
+
+    /**
+     * Find the LCA of two vertexes
+     * @param u a vertex of pair to find the lca
+     * @param v a vertex of pair to find the lca
+     * @return the lca of u and v
+     */
+    pub fn query(&self, mut u: usize, mut v: usize) -> usize {
+        if self.depth[u] > self.depth[v] {
+            std::mem::swap(&mut u, &mut v);
+        }
+
+        let height = self.depth[v] - self.depth[u];
+
+        if height != 0 {
+            let lg = 31 - height.leading_zeros() as usize;
+
+            for i in 0..=lg {
+                if (height >> i) & 1 == 1 {
+                    v = self.st[v][i];
+                }
+            }
+        }
+
+        if u == v {
+            return u;
+        }
+
+        let lg = 31 - self.depth[u].leading_zeros() as usize;
+
+        for i in (0..=lg).rev() {
+            if self.st[u][i] == self.st[v][i] {
+                continue;
+            }
+
+            u = self.st[u][i];
+            v = self.st[v][i];
+        }
+
+        self.st[u][0]
+    }
+
+    /**
+     * Calculate the distance between two nodes
+     * @param u the vertex of the pair to find the distance
+     * @param v the vertex of the pair to find the distance
+     * @return the distance between u and v
+     */
+    pub fn distance(&self, u: usize, v: usize) -> i64 {
+        let x = self.query(u, v);
+        self.weight[u] + self.weight[v] - 2 * self.weight[x]
+    }
+}

competitive_programming/src/graphs/mod.rs

@@ -1,3 +1,5 @@
+mod binary_lifting;
 mod tree_diameter;
 
+pub use binary_lifting::LCA;
 pub use tree_diameter::diameter;

competitive_programming/tests/graphs/binary_lifting.rs

@@ -0,0 +1,118 @@
+#[cfg(test)]
+mod lca_tests {
+    use competitive_programming::graphs::LCA;
+
+    #[test]
+    fn simple_path_test() {
+        const N: usize = 10;
+
+        let mut lca = LCA::new(N);
+
+        for i in 1..N {
+            lca.add_edge(i - 1, i, 1);
+        }
+
+        lca.build(0);
+
+        for u in 0..N {
+            for v in 0..N {
+                let distance = lca.distance(u, v);
+                let vertex = lca.query(u, v);
+                let expected_distance = u.abs_diff(v);
+                let expected_vertex = u.min(v);
+
+                assert_eq!(distance, expected_distance as i64);
+                assert_eq!(vertex, expected_vertex);
+            }
+        }
+    }
+
+    #[test]
+    fn star_test() {
+        const N: usize = 10;
+
+        let mut lca = LCA::new(N);
+
+        for i in 1..N {
+            lca.add_edge(0, i, 1);
+        }
+
+        lca.build(0);
+
+        for i in 1..N {
+            let distance = lca.distance(0, i);
+            let expected_distance = 1;
+
+            assert_eq!(distance, expected_distance);
+            assert_eq!(lca.query(0, i), 0);
+        }
+
+        for u in 1..N {
+            for v in 1..N {
+                if u == v {
+                    continue;
+                }
+
+                let distance = lca.distance(u, v);
+                let expected_distance = 2;
+
+                assert_eq!(distance, expected_distance);
+                assert_eq!(lca.query(u, v), 0);
+            }
+        }
+    }
+
+    #[test]
+    #[allow(unused_assignments)]
+    fn double_star_test() {
+        const N: usize = 10;
+
+        let mut lca = LCA::new(N);
+
+        lca.add_edge(0, 1, 7);
+
+        for i in 2..N {
+            lca.add_edge(i % 2, i, 7);
+        }
+
+        lca.build(0);
+
+        for u in 0..2 {
+            for v in 0..N {
+                if u == v {
+                    assert_eq!(lca.distance(u, v), 0);
+                    assert_eq!(lca.query(u, v), u);
+                } else if u == v ^ 1 {
+                    assert_eq!(lca.distance(u, v), 7);
+                    assert_eq!(lca.query(u, v), 0);
+                } else if u % 2 == v % 2 {
+                    assert_eq!(lca.distance(u, v), 7);
+                    assert_eq!(lca.query(u, v), u);
+                } else {
+                    assert_eq!(lca.distance(u, v), 14);
+                    assert_eq!(lca.query(u, v), 0);
+                }
+            }
+        }
+
+        for u in 3..N {
+            for v in (u + 1)..N {
+                let distance = lca.distance(u, v);
+                let vertex = lca.query(u, v);
+
+                let mut expected_distance = 0;
+                let mut expected_vertex = 0;
+
+                if u % 2 == v % 2 {
+                    expected_distance = 14;
+                    expected_vertex = u % 2;
+                } else {
+                    expected_distance = 21;
+                }
+
+                assert_eq!(distance, expected_distance);
+                assert_eq!(vertex, expected_vertex);
+            }
+        }
+    }
+}

competitive_programming/tests/graphs/mod.rs

@@ -1 +1,2 @@
+pub mod binary_lifting;
 pub mod tree_diameter;