haar_lib/tree/
lca.rs

1//! 最小共通祖先
2
3use crate::tree::*;
4
5/// ダブリングによる最小共通祖先
6pub struct DoublingLCA {
7    log2n: usize,
8    parent: Vec<Vec<Option<usize>>>,
9    depth: Vec<usize>,
10}
11
12impl DoublingLCA {
13    /// **Time complexity** $O(n \log n)$
14    ///
15    /// **Space complexity** $O(n \log n)$
16    pub fn new<E: TreeEdgeTrait>(tree: &Tree<E>, root: usize) -> Self {
17        let n = tree.len();
18        let log2n = n.next_power_of_two().trailing_zeros() as usize + 1;
19        let mut this = Self {
20            log2n,
21            parent: vec![vec![None; log2n]; n],
22            depth: vec![0; n],
23        };
24
25        let mut stack = vec![];
26        stack.push((root, None, 0));
27
28        while let Some((cur, par, d)) = stack.pop() {
29            this.parent[cur][0] = par;
30            this.depth[cur] = d;
31
32            tree.nodes[cur]
33                .neighbors()
34                .filter(|e| par.is_none_or(|p| p != e.to()))
35                .for_each(|e| stack.push((e.to(), Some(cur), d + 1)));
36        }
37
38        for k in 0..log2n - 1 {
39            for v in 0..n {
40                match this.parent[v][k] {
41                    Some(p) => this.parent[v][k + 1] = this.parent[p][k],
42                    None => this.parent[v][k + 1] = None,
43                }
44            }
45        }
46
47        this
48    }
49
50    /// `a`の`n`個上の祖先を求める。
51    ///
52    /// **Time complexity** $O(\log n)$
53    pub fn ancestor(&self, mut a: usize, mut n: usize) -> Option<usize> {
54        while n != 0 {
55            let m1 = usize::BITS as usize - n.leading_zeros() as usize - 1;
56            if let Some(&Some(b)) = self.parent[a].get(m1) {
57                a = b;
58                n ^= 1 << m1;
59            } else {
60                return None;
61            }
62        }
63
64        Some(a)
65    }
66
67    /// `a`と`b`の最小共通祖先を求める。
68    ///
69    /// **Time complexity** $O(\log n)$
70    pub fn lca(&self, mut a: usize, mut b: usize) -> usize {
71        if self.depth[a] >= self.depth[b] {
72            std::mem::swap(&mut a, &mut b);
73        }
74        for k in 0..self.log2n {
75            if (((self.depth[b] - self.depth[a]) >> k) & 1) != 0 {
76                b = self.parent[b][k].unwrap();
77            }
78        }
79        if a == b {
80            return a;
81        }
82
83        for k in (0..self.log2n).rev() {
84            if self.parent[a][k] != self.parent[b][k] {
85                a = self.parent[a][k].unwrap();
86                b = self.parent[b][k].unwrap();
87            }
88        }
89
90        self.parent[a][0].unwrap()
91    }
92
93    /// s-t最短パス上で、sから見てd番目の頂点を返す。
94    ///
95    /// **Time complexity** $O(\log n)$
96    pub fn jump(&self, s: usize, t: usize, d: usize) -> Option<usize> {
97        let a = self.lca(s, t);
98        if self.depth[s] - self.depth[a] >= d {
99            self.ancestor(s, d)
100        } else if self.depth[s] + self.depth[t] - self.depth[a] * 2 >= d {
101            self.ancestor(t, self.depth[s] + self.depth[t] - self.depth[a] * 2 - d)
102        } else {
103            None
104        }
105    }
106}
haar_lib/tree/lca.rs

haar_lib/tree/
lca.rs
