competitive_programming_rs/graph/
lca.rs

1pub mod lca {
2    const MAX_PARENT: usize = 1 << 50;
3    pub struct LowestCommonAncestor {
4        parent: Vec<Vec<usize>>,
5        depth: Vec<usize>,
6        log_v: usize,
7    }
8
9    impl LowestCommonAncestor {
10        pub fn new(graph: &[Vec<usize>]) -> Self {
11            let num_v = graph.len();
12            let root = 0;
13            let mut depth = vec![0; num_v];
14
15            let mut log_v = 1;
16            let mut i = 1;
17            while i <= num_v {
18                i *= 2;
19                log_v += 1;
20            }
21            let mut parent: Vec<Vec<usize>> = vec![vec![0; num_v]; log_v];
22
23            let mut depth_vis = vec![false; num_v];
24            let mut stack = ::std::collections::VecDeque::new();
25            stack.push_front(root);
26            parent[0][root] = MAX_PARENT;
27            depth[root] = 0;
28            depth_vis[root] = true;
29            while !stack.is_empty() {
30                let v = stack.pop_front().unwrap();
31                stack.push_front(v);
32                for &u in &graph[v] {
33                    if depth_vis[u] {
34                        continue;
35                    }
36                    parent[0][u] = v;
37                    depth[u] = depth[v] + 1;
38                    depth_vis[u] = true;
39                    stack.push_front(u);
40                }
41
42                let head = stack.pop_front().unwrap();
43                if head != v {
44                    stack.push_front(head);
45                }
46            }
47
48            for k in 0..(log_v - 1) {
49                for u in 0..num_v {
50                    parent[k + 1][u] = if parent[k][u] == MAX_PARENT {
51                        MAX_PARENT
52                    } else {
53                        parent[k][parent[k][u]]
54                    };
55                }
56            }
57
58            LowestCommonAncestor {
59                parent,
60                depth,
61                log_v,
62            }
63        }
64
65        pub fn get_lca(&self, u: usize, v: usize) -> usize {
66            let (mut u, mut v) = if self.depth[u] <= self.depth[v] {
67                (u, v)
68            } else {
69                (v, u)
70            };
71            for k in 0..self.log_v {
72                if ((self.depth[v] - self.depth[u]) & (1 << k)) != 0 {
73                    v = self.parent[k][v];
74                }
75            }
76            if u == v {
77                return u;
78            }
79
80            for k in (0..self.log_v).rev() {
81                if self.parent[k][u] != self.parent[k][v] {
82                    u = self.parent[k][u];
83                    v = self.parent[k][v];
84                }
85            }
86            self.parent[0][u]
87        }
88
89        pub fn get_dist(&self, u: usize, v: usize) -> usize {
90            let lca = self.get_lca(u, v);
91            self.depth[u] + self.depth[v] - self.depth[lca] * 2
92        }
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::lca::*;
99    use crate::utils::test_helper::Tester;
100
101    #[test]
102    fn solve_grl_5_c() {
103        let tester = Tester::new("./assets/GRL_5_C/in/", "./assets/GRL_5_C/out/");
104        tester.test_solution(|sc| {
105            let n: usize = sc.read();
106            let mut graph = (0..n).map(|_| Vec::new()).collect::<Vec<_>>();
107            for i in 0..n {
108                let k = sc.read();
109                for _ in 0..k {
110                    let c: usize = sc.read();
111                    graph[i].push(c);
112                    graph[c].push(i);
113                }
114            }
115
116            let lca = LowestCommonAncestor::new(&graph);
117
118            let q: usize = sc.read();
119            for _ in 0..q {
120                let u: usize = sc.read();
121                let v: usize = sc.read();
122                let ans = lca.get_lca(u, v);
123                sc.write(format!("{}\n", ans));
124            }
125        });
126    }
127}
competitive_programming_rs/graph/lca.rs

competitive_programming_rs/graph/
lca.rs
