competitive_library/graph/
heavy_light_decomposition.rs

1//! HL 分解
2
3#[derive(Debug, Clone)]
4pub struct HeavyLightDecomposition {
5    /// 根
6    root: usize,
7    /// 各頂点の親
8    parent: Vec<usize>,
9    /// e\[i]\[j] は i を親に持つ
10    e: Vec<Vec<usize>>,
11    /// 各頂点から見た子要素数
12    child_count: Vec<usize>,
13    /// 各頂点の深さ
14    depths: Vec<usize>,
15    pre: Vec<usize>,
16    /// HLD
17    hld: Vec<usize>,
18    /// 列の先頭
19    head: Vec<usize>,
20}
21
22impl HeavyLightDecomposition {
23    #[inline]
24    pub fn new(root: usize, parent: &[usize]) -> Self {
25        let mut e = vec![vec![]; parent.len()];
26        for (i, &v) in parent.iter().enumerate().filter(|&x| x.0 != *x.1) {
27            e[v].push(i);
28        }
29
30        Self {
31            root,
32            parent: parent.to_vec(),
33            e,
34            child_count: vec![0; parent.len()],
35            depths: vec![0; parent.len()],
36            pre: vec![0; parent.len()],
37            hld: vec![],
38            head: (0..parent.len()).collect(),
39        }
40    }
41
42    /// 分解
43    #[inline]
44    pub fn decompose(&mut self) -> Vec<usize> {
45        let init = self.root;
46        self.count_node(init);
47        self.count_depth(init);
48        self.decompose_inner_root(init);
49
50        self.hld.clone()
51    }
52
53    #[inline]
54    /// 根から分解
55    fn decompose_inner_root(&mut self, v: usize) {
56        self.decompose_inner(v, v);
57    }
58
59    /// 分解用の内部関数
60    #[inline]
61    fn decompose_inner(&mut self, v: usize, h: usize) {
62        self.pre[v] = self.hld.len();
63        self.hld.push(v);
64        self.head[v] = h;
65
66        if self.e[v].is_empty() {
67            return;
68        }
69        let index = self.e[v]
70            .iter()
71            .enumerate()
72            .max_by_key(|&(_, &y)| self.child_count[y])
73            .unwrap()
74            .0;
75        self.decompose_inner(self.e[v][index], h);
76
77        for i in (0..self.e[v].len()).filter(|&i| i != index) {
78            self.decompose_inner_root(self.e[v][i]);
79        }
80    }
81
82    /// 子要素のカウント
83    #[inline]
84    fn count_node(&mut self, index: usize) -> usize {
85        if self.child_count[index] != 0 {
86            return self.child_count[index];
87        }
88        self.child_count[index] = 1;
89        for i in 0..self.e[index].len() {
90            self.child_count[index] += self.count_node(self.e[index][i]);
91        }
92        self.child_count[index]
93    }
94
95    /// 深さのカウント
96    #[inline]
97    fn count_depth(&mut self, index: usize) -> usize {
98        if self.depths[index] != 0 {
99            return self.depths[index];
100        }
101        if self.parent[index] == index {
102            return 0;
103        }
104        self.depths[index] = self.count_depth(self.parent[index]) + 1;
105        self.depths[index]
106    }
107
108    /// HLD 配列の区間を返す
109    #[inline]
110    pub fn query(&mut self, mut u: usize, mut v: usize) -> Vec<(usize, usize)> {
111        debug_assert!(!self.hld.is_empty());
112
113        let mut ret = vec![];
114        while self.head[u] != self.head[v] {
115            if self.count_depth(self.head[u]) <= self.count_depth(self.head[v]) {
116                ret.push((self.pre[self.head[v]], self.pre[v]));
117                v = self.parent[self.head[v]];
118            } else {
119                ret.push((self.pre[self.head[u]], self.pre[u]));
120                u = self.parent[self.head[u]];
121            }
122        }
123        ret.push(if self.pre[u] < self.pre[v] {
124            (self.pre[u], self.pre[v])
125        } else {
126            (self.pre[v], self.pre[u])
127        });
128        ret
129    }
130
131    /// LCA を求める
132    #[inline]
133    pub fn get_lca(&mut self, u: usize, v: usize) -> Option<usize> {
134        let common_range = *self.query(u, v).last()?;
135        Some(self.hld[common_range.0])
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    #[test]
143    fn test_hld() {
144        let v = vec![0, 0, 1, 2, 2, 1, 0, 6, 7, 7, 0, 10];
145
146        let mut hld = HeavyLightDecomposition::new(0, &v);
147        let h = hld.decompose();
148        dbg!(&h);
149
150        use std::collections::HashSet;
151        let mut set = HashSet::new();
152        for (f, t) in hld.query(4, 9) {
153            for &i in h.iter().take(t + 1).skip(f) {
154                set.insert(i);
155            }
156        }
157        let ans_set = [4_usize, 2, 1, 0, 6, 7, 9].iter().cloned().collect();
158
159        assert_eq!(set, ans_set);
160    }
161    #[test]
162    fn test_lca() {
163        let v = vec![0, 0, 0, 2, 2];
164
165        let mut hld = HeavyLightDecomposition::new(0, &v);
166        let h = hld.decompose();
167        dbg!(&h);
168        for &(u, v, ans) in [
169            (0, 0, 0),
170            (0, 1, 0),
171            (0, 2, 0),
172            (0, 3, 0),
173            (0, 4, 0),
174            (1, 1, 1),
175            (1, 2, 0),
176            (1, 3, 0),
177            (1, 4, 0),
178            (2, 2, 2),
179            (2, 3, 2),
180            (2, 4, 2),
181            (3, 3, 3),
182            (3, 4, 2),
183            (4, 4, 4),
184        ]
185        .iter()
186        {
187            assert_eq!(hld.get_lca(u, v).unwrap(), ans, "{} {}", u, v);
188        }
189    }
190}
competitive_library/graph/heavy_light_decomposition.rs

competitive_library/graph/
heavy_light_decomposition.rs
