cp_library_rs/graph/
lca_doubling.rs

1//! ダブリングにより、最小共通祖先を求める
2
3use crate::utils::consts::Infinity;
4
5type Graph = Vec<Vec<usize>>;
6
7/// LCA
8/// - 最小共通祖先を求めるクエリに答える
9pub struct LCA {
10    root: usize,
11    logn: usize,
12    double: Vec<Vec<usize>>,
13    depth: Vec<usize>,
14}
15
16impl LCA {
17    /// `root`を根に持つ木`tree`で、初期化を行う
18    pub fn new(tree: &Graph, root: usize) -> Self {
19        let n = tree.len(); // グラフの頂点数
20        let logn = n.next_power_of_two().trailing_zeros() as usize;
21        let mut double = vec![vec![0; n]; logn]; // ダブリング配列
22        let mut depth = vec![usize::infinity(); n]; // 頂点の根からの距離
23        depth[0] = 0;
24        Self::dfs(root, &mut double[0], &mut depth, tree);
25
26        // ダブリング
27        for i in 1..logn {
28            for j in 0..n {
29                double[i][j] = double[i - 1][double[i - 1][j]];
30            }
31        }
32
33        Self {
34            root,
35            logn,
36            double,
37            depth,
38        }
39    }
40
41    fn dfs(u: usize, par: &mut Vec<usize>, depth: &mut Vec<usize>, tree: &Graph) {
42        for &v in &tree[u] {
43            if depth[v] != usize::infinity() {
44                continue;
45            }
46            depth[v] = depth[u] + 1;
47            par[v] = u;
48            Self::dfs(v, par, depth, tree);
49        }
50    }
51
52    /// 頂点 `v` の深さを求める
53    pub fn depth(&self, v: usize) -> usize {
54        self.depth[v]
55    }
56
57    /// 頂点 `v` の親を求める
58    pub fn parent(&self, v: usize) -> Option<usize> {
59        (v != self.root).then_some(self.double[0][v])
60    }
61
62    /// 頂点 `u`,`v` の最小共通祖先を求める
63    pub fn lca(&self, mut u: usize, mut v: usize) -> usize {
64        // 常にuを深くする
65        if self.depth[u] < self.depth[v] {
66            (u, v) = (v, u);
67        }
68
69        // LCAまでの距離を同じにする
70        for k in 0..self.logn {
71            if ((self.depth[u] - self.depth[v]) >> k) & 1 == 1 {
72                u = self.double[k][u];
73            }
74        }
75
76        if u == v {
77            return u;
78        }
79
80        // 二分探索
81        for k in (0..self.logn).rev() {
82            if self.double[k][u] != self.double[k][v] {
83                u = self.double[k][u];
84                v = self.double[k][v];
85            }
86        }
87
88        self.double[0][u]
89    }
90
91    /// 頂点 `u`,`v` の距離を求める
92    pub fn dist(&self, u: usize, v: usize) -> usize {
93        let o = self.lca(u, v);
94        (self.depth[u] - self.depth[o]) + (self.depth[v] - self.depth[o])
95    }
96
97    /// 頂点 `v` の k 個上の祖先を求める
98    pub fn kth_ancestor(&self, mut v: usize, k: usize) -> Option<usize> {
99        if k > self.depth(v) {
100            return None;
101        }
102        for i in 0..self.logn {
103            if (k >> i) & 1 == 1 {
104                v = self.double[i][v];
105            }
106        }
107        Some(v)
108    }
109
110    /// 頂点 `u` から `v` へ向かうパス上の `k` 個目の頂点を求める
111    pub fn kth_on_path(&self, u: usize, v: usize, k: usize) -> Option<usize> {
112        let o = self.lca(u, v);
113        let dist_u_o = self.depth[u] - self.depth[o];
114        if k <= dist_u_o {
115            self.kth_ancestor(u, k)
116        } else {
117            let dist_v_o = self.depth[v] - self.depth[o];
118            (dist_v_o + dist_u_o)
119                .checked_sub(k)
120                .and_then(|l| self.kth_ancestor(v, l))
121        }
122    }
123}
cp_library_rs/graph/lca_doubling.rs

cp_library_rs/graph/
lca_doubling.rs
