competitive_library/graph/
lowest_common_ancestor_rmq.rs

1//! RMQ を使用しLCAを求める
2
3use crate::graph::euler_tour::euler_tour;
4use crate::structure::sparse_table::{Band, SparseTable};
5
6/// 深さ及び頂点番号を保持しの最小値を返す
7struct MinDepth {}
8impl Band for MinDepth {
9    /// (vertex, depth)
10    type T = (i32, i32);
11
12    /// Min depth
13    fn operate(a: &Self::T, b: &Self::T) -> Self::T {
14        if a.1 < b.1 { *a } else { *b }
15    }
16}
17/// LCA
18pub struct LowestCommonAncestor {
19    st: SparseTable<MinDepth>,
20    first_look: Vec<usize>,
21}
22
23impl LowestCommonAncestor {
24    /// ## Arguments
25    /// * `e` - 根付き木の辺集合。隣接リスト
26    /// * `root` - 根付き木の根
27    #[inline]
28    pub fn new(e: &[Vec<usize>], root: usize) -> Self {
29        let (tour, first_look, depths) = euler_tour(e, root);
30        let v = tour
31            .iter()
32            .map(|&x| (x as i32, depths[x] as i32))
33            .collect::<Vec<_>>();
34        let st = SparseTable::new(&v);
35
36        LowestCommonAncestor { st, first_look }
37    }
38
39    /// 対象の2頂点のLCAを取得する
40    ///
41    /// ## Arguments
42    /// * `u` - 頂点番号
43    /// * `v` - 頂点番号
44    ///
45    /// ## Returns
46    /// LCA
47    #[inline]
48    pub fn get_lca(&self, u: usize, v: usize) -> usize {
49        if u == v {
50            return u;
51        }
52        let range = if self.first_look[u] < self.first_look[v] {
53            self.first_look[u]..self.first_look[v]
54        } else {
55            self.first_look[v]..self.first_look[u]
56        };
57        self.st.fold(range).0 as usize
58    }
59}
60
61#[cfg(test)]
62mod tests {
63
64    use super::*;
65    #[test]
66    fn test_lca() {
67        let n = 5;
68        let mut e = vec![vec![]; n];
69        for (i, &v) in [0, 0, 2, 2].iter().enumerate() {
70            e[v].push(i + 1);
71        }
72
73        let lca = LowestCommonAncestor::new(&e, 0);
74        for &(u, v, ans) in [
75            (0, 0, 0),
76            (0, 1, 0),
77            (0, 2, 0),
78            (0, 3, 0),
79            (0, 4, 0),
80            (1, 1, 1),
81            (1, 2, 0),
82            (1, 3, 0),
83            (1, 4, 0),
84            (2, 2, 2),
85            (2, 3, 2),
86            (2, 4, 2),
87            (3, 3, 3),
88            (3, 4, 2),
89            (4, 4, 4),
90        ]
91        .iter()
92        {
93            assert_eq!(lca.get_lca(u, v), ans);
94        }
95    }
96}