1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79

use std::collections::HashMap;

use lowest_common_ancestor::LowestCommonAncestor;

/// [auxiliary tree](https://noshi91.github.io/algorithm-encyclopedia/auxiliary-tree) です。
///
/// # 計算量
///
/// hash map のコストは無視する。
///
/// グラフの頂点数を n, `nodes.len()` を k として、O(klogn + klogk)
///
/// # 引数
///
/// * `nodes`： {0, 1, ..., n - 1} の部分集合
/// * `inv_ord`: pre-order (行きがけ順) の列の添字と値を反対にしたもの
///     * 頂点 `i` は行きがけ順で `inv_ord[i]` 番目に現われる
/// * `lca`: 2頂点の LCA を得る `.get(u, v)` を実装した構造体
///
/// # 返り値
///
/// `(root, g)`
///
/// * ` g.contains_key(&i)`: 頂点 `i` が圧縮後の木に含まれて、子のリストが `g[&i]` である
/// * `!g.contains_key(&i)`: 頂点 `i` が圧縮後の木に含まれない
pub fn auxiliary_tree(
    nodes: &[usize],
    inv_ord: &[usize],
    lca: &LowestCommonAncestor, // trait にする？
) -> (usize, HashMap<usize, Vec<usize>>) {
    // https://smijake3.hatenablog.com/entry/2019/09/15/200200

    assert!(!nodes.is_empty());

    // nodes.len() < 2 だと .windows(2) が空になるので場合分け
    if nodes.len() == 1 {
        return (nodes[0], HashMap::from([(nodes[0], vec![])]));
    }

    let mut nodes = nodes.to_vec();
    nodes.sort_by_key(|&i| inv_ord[i]);

    let lca_nodes = nodes
        .windows(2)
        .map(|w| lca.get(w[0], w[1]))
        .collect::<Vec<_>>();
    nodes.extend(lca_nodes);
    nodes.sort_by_key(|&i| inv_ord[i]);
    nodes.dedup();

    let mut h = HashMap::<_, Vec<_>>::new();
    for w in nodes.windows(2) {
        // stack 使わずにこれでよさそう
        let x = lca.get(w[0], w[1]);
        h.entry(x).or_insert_with(Vec::new).push(w[1]);
        assert!(!h.contains_key(&w[1]));
        h.insert(w[1], vec![]);
    }
    (nodes[0], h)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_line() {
        //                  *         *
        // 0 (root) -- 1 -- 2 -- 3 -- 4
        assert_eq!(
            auxiliary_tree(
                &[2, 4],
                &[0, 1, 2, 3, 4],
                &LowestCommonAncestor::new(5, 0, &[(0, 1), (1, 2), (2, 3), (3, 4)])
            ),
            (2, HashMap::from([(2, vec![4]), (4, vec![])]))
        );
    }
}