cp_library_rs/graph/
hld.rs

1//! ## HL分解（重軽分解）
2
3use crate::utils::consts::Infinity;
4
5/// HL分解
6pub struct HLD {
7    /// 頂点数
8    pub N: usize,
9    /// 根
10    pub root: usize,
11    /// グラフ
12    pub G: Vec<Vec<usize>>,
13    /// 親頂点
14    pub parent: Vec<usize>,
15    /// subtree_size[i] := `i`を根とする部分木のサイズ
16    pub subtree_size: Vec<usize>,
17    /// 行きがけ順での番号
18    pub in_: Vec<usize>,
19    /// 帰りがけ順での番号
20    pub out: Vec<usize>,
21    /// heavy pathの端点
22    pub head: Vec<usize>,
23}
24
25impl HLD {
26    /// `N`頂点の木を初期化する
27    pub fn new(N: usize) -> Self {
28        Self {
29            N,
30            root: usize::infinity(),
31            G: vec![vec![]; N],
32            parent: vec![usize::infinity(); N],
33            subtree_size: vec![usize::infinity(); N],
34            in_: vec![usize::infinity(); N],
35            out: vec![usize::infinity(); N],
36            head: vec![usize::infinity(); N],
37        }
38    }
39
40    /// 辺`(u,v)`を追加する
41    pub fn add_edge(&mut self, u: usize, v: usize) {
42        self.G[u].push(v);
43        self.G[v].push(u);
44    }
45
46    /// 頂点`root`を根としてHL分解をする
47    pub fn decompose(&mut self, root: usize) {
48        self.root = root;
49
50        // heavy childの計算
51        self.set_heavy_child(usize::infinity(), root);
52
53        // heavy pathの計算
54        self.head[root] = root;
55        self.build_heavy_path(usize::infinity(), root, &mut 0);
56    }
57
58    /// 再帰的にheavy childを計算し，
59    /// heavy childが`G[u][0]`にくるように設定する．
60    ///
61    /// （これにより，行きがけ順の走査で自然にheavy pathがえられる）
62    fn set_heavy_child(&mut self, p: usize, u: usize) {
63        self.parent[u] = p;
64        self.subtree_size[u] = 1;
65
66        for i in 0..self.G[u].len() {
67            let v = self.G[u][i];
68            if v == p {
69                continue;
70            }
71            // 再帰的に計算
72            self.set_heavy_child(u, v);
73
74            // 部分木のサイズを足す
75            self.subtree_size[u] += self.subtree_size[v];
76
77            // G[u][0]にheavy childがくるようにswap
78            let v_0 = self.G[u][0];
79            if v_0 == p || self.subtree_size[v] > self.subtree_size[v_0] {
80                self.G[u].swap(i, 0);
81            }
82        }
83    }
84
85    /// 行きがけ順に走査し，heavy pathの列を構築する
86    fn build_heavy_path(&mut self, p: usize, u: usize, id: &mut usize) {
87        self.in_[u] = *id;
88        *id += 1;
89
90        for i in 0..self.G[u].len() {
91            let v = self.G[u][i];
92            if v == p {
93                continue;
94            }
95            self.head[v] = if i == 0 {
96                // 自分がheavy childの場合
97                self.head[u]
98            } else {
99                v
100            };
101
102            self.build_heavy_path(u, v, id);
103        }
104
105        self.out[u] = *id;
106    }
107
108    /// 頂点`v`の配列上でのインデックス
109    #[inline]
110    pub fn get_id(&self, v: usize) -> usize {
111        self.in_[v]
112    }
113
114    /// 2頂点`u,v`の最小共通祖先 (LCA) を求める
115    pub fn get_lca(&self, mut u: usize, mut v: usize) -> usize {
116        let mut pu = self.head[u];
117        let mut pv = self.head[v];
118
119        while self.head[u] != self.head[v] {
120            if self.in_[pu] > self.in_[pv] {
121                u = self.parent[pu];
122                pu = self.head[u];
123            } else {
124                v = self.parent[pv];
125                pv = self.head[v];
126            }
127        }
128
129        if self.in_[u] <= self.in_[v] {
130            u
131        } else {
132            v
133        }
134    }
135
136    /// `u`を根とする部分木の値を集約する
137    ///
138    /// （モノイド`M`が可環であるときに定義される）
139    ///
140    /// **戻り値**
141    /// - `(usize, usize)` : (左, 右)
142    #[inline]
143    pub fn get_subtree<T, F>(&self, u: usize) -> (usize, usize) {
144        (self.in_[u], self.out[u])
145    }
146
147    /// 頂点`u,v`間のパス上のパス断片を順に返すイテレータを返す．
148    ///
149    /// **戻り値 (Item)**
150    /// - `i (usize)` : 最も根側の頂点
151    /// - `j (usize)` : 最も葉側の頂点
152    /// - `last (bool)` : 最後のpath segmentであるか
153    /// - `rev (bool)` : 取得するpathの向きに対してpath segmentが逆方向かどうか
154    pub fn get_path(&self, u: usize, v: usize) -> PathSegment<'_> {
155        PathSegment {
156            hld: self,
157            from: u,
158            to: v,
159            exhasted: false,
160        }
161    }
162}
163
164/// パス断片を返すイテレータ
165///
166/// **戻り値 (next)**
167/// - `i (usize)` : 最も根側の頂点
168/// - `j (usize)` : 最も葉側の頂点
169/// - `last (bool)` : 最後のpath segmentであるか
170/// - `rev (bool)` : 取得するpathの向きに対してpath segmentが逆方向かどうか
171///
172/// **参考**
173/// - <https://ngtkana.hatenablog.com/entry/2024/06/24/200138>
174pub struct PathSegment<'a> {
175    hld: &'a HLD,
176    from: usize,
177    to: usize,
178    exhasted: bool,
179}
180
181impl Iterator for PathSegment<'_> {
182    type Item = (usize, usize, bool, bool);
183
184    fn next(&mut self) -> Option<Self::Item> {
185        let Self {
186            hld,
187            from,
188            to,
189            exhasted,
190        } = *self;
191
192        if exhasted {
193            return None;
194        }
195
196        let HLD {
197            in_, head, parent, ..
198        } = hld;
199
200        Some(if head[from] == head[to] {
201            self.exhasted = true;
202            if in_[from] < in_[to] {
203                (from, to, true, false)
204            } else {
205                (to, from, true, true)
206            }
207        } else if in_[from] < in_[to] {
208            self.to = parent[head[to]];
209            (head[to], to, false, false)
210        } else {
211            self.from = parent[head[from]];
212            (head[from], from, false, true)
213        })
214    }
215}
cp_library_rs/graph/hld.rs

cp_library_rs/graph/
hld.rs
