auxiliary_tree/
lib.rs

1use std::collections::HashMap;
2
3use lowest_common_ancestor::LowestCommonAncestor;
4
5/// [auxiliary tree](https://noshi91.github.io/algorithm-encyclopedia/auxiliary-tree) です。
6///
7/// # 計算量
8///
9/// hash map のコストは無視する。
10///
11/// グラフの頂点数を n, `nodes.len()` を k として、O(klogn + klogk)
12///
13/// # 引数
14///
15/// * `nodes`： {0, 1, ..., n - 1} の部分集合
16/// * `inv_ord`: pre-order (行きがけ順) の列の添字と値を反対にしたもの
17///     * 頂点 `i` は行きがけ順で `inv_ord[i]` 番目に現われる
18/// * `lca`: 2頂点の LCA を得る `.get(u, v)` を実装した構造体
19///
20/// # 返り値
21///
22/// `(root, g)`
23///
24/// * ` g.contains_key(&i)`: 頂点 `i` が圧縮後の木に含まれて、子のリストが `g[&i]` である
25/// * `!g.contains_key(&i)`: 頂点 `i` が圧縮後の木に含まれない
26pub fn auxiliary_tree(
27    nodes: &[usize],
28    inv_ord: &[usize],
29    lca: &LowestCommonAncestor, // trait にする？
30) -> (usize, HashMap<usize, Vec<usize>>) {
31    // https://smijake3.hatenablog.com/entry/2019/09/15/200200
32
33    assert!(!nodes.is_empty());
34
35    // nodes.len() < 2 だと .windows(2) が空になるので場合分け
36    if nodes.len() == 1 {
37        return (nodes[0], HashMap::from([(nodes[0], vec![])]));
38    }
39
40    let mut nodes = nodes.to_vec();
41    nodes.sort_by_key(|&i| inv_ord[i]);
42
43    let lca_nodes = nodes
44        .windows(2)
45        .map(|w| lca.get(w[0], w[1]))
46        .collect::<Vec<_>>();
47    nodes.extend(lca_nodes);
48    nodes.sort_by_key(|&i| inv_ord[i]);
49    nodes.dedup();
50
51    let mut h = HashMap::<_, Vec<_>>::new();
52    for w in nodes.windows(2) {
53        // stack 使わずにこれでよさそう
54        let x = lca.get(w[0], w[1]);
55        h.entry(x).or_insert_with(Vec::new).push(w[1]);
56        assert!(!h.contains_key(&w[1]));
57        h.insert(w[1], vec![]);
58    }
59    (nodes[0], h)
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65
66    #[test]
67    fn test_line() {
68        //                  *         *
69        // 0 (root) -- 1 -- 2 -- 3 -- 4
70        assert_eq!(
71            auxiliary_tree(
72                &[2, 4],
73                &[0, 1, 2, 3, 4],
74                &LowestCommonAncestor::new(5, 0, &[(0, 1), (1, 2), (2, 3), (3, 4)])
75            ),
76            (2, HashMap::from([(2, vec![4]), (4, vec![])]))
77        );
78    }
79}
auxiliary_tree/lib.rs

auxiliary_tree/
lib.rs
