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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
//! ## UnionFind木
//!
//! モノイドを乗せるUnionFind木.

use std::{collections::HashMap, fmt::Debug, mem};

use crate::{algebraic_structure::commutative::CommutativeMonoid, utils::consts::NEG1};

/// UnionFind木
pub type UnionFind = UnionFindMonoid<()>;

/// UnionFind木(モノイド)
pub struct UnionFindMonoid<M: CommutativeMonoid> {
    /// 要素数
    n: usize,
    /// 親の番号を格納する配列
    parent: Vec<usize>,
    /// 値
    value: Vec<Option<M::Val>>,
    /// 連結成分の個数
    count: usize,
}

impl<M: CommutativeMonoid> UnionFindMonoid<M> {
    /// 根を求める
    pub fn root(&mut self, mut x: usize) -> usize {
        // 根を探索
        let mut root = x;
        while self.parent[root] < self.n {
            root = self.parent[root];
        }
        // 経路圧縮
        while self.parent[x] < self.n {
            x = mem::replace(&mut self.parent[x], root);
        }
        root
    }

    /// ノード`x`が属する集合の値を取得
    pub fn value(&mut self, x: usize) -> &M::Val {
        let root = self.root(x);
        self.value[root].as_ref().unwrap()
    }

    /// 同一の集合に所属するか判定
    pub fn is_same(&mut self, x: usize, y: usize) -> bool {
        self.root(x) == self.root(y)
    }

    /// 集合`x,y`を併合する.
    ///
    /// **戻り値**
    /// - すでに併合済みだった場合`None`,そうでない場合親となった要素の番号を返す
    pub fn unite(&mut self, x: usize, y: usize) -> Option<usize> {
        let mut parent = self.root(x);
        let mut child = self.root(y);

        if parent == child {
            return None;
        }

        // 要素数が大きい方を親にすることで、高さを均等に保つ
        if self.parent[parent] > self.parent[child] {
            (parent, child) = (child, parent);
        }

        self.parent[parent] = self.parent[parent].wrapping_add(self.parent[child]);
        self.parent[child] = parent;
        self.count -= 1;

        // 値のマージ
        let child_val = self.value[child].take();
        let parent_val = self.value[parent].take();
        self.value[parent] = child_val.zip(parent_val).map(|(c, p)| M::op(&c, &p));

        Some(parent)
    }

    /// 連結成分の大きさを求める
    pub fn get_size(&mut self, x: usize) -> usize {
        let root = self.root(x);
        self.parent[root].wrapping_neg()
    }

    /// 連結成分の数を返す
    pub fn group_count(&self) -> usize {
        self.count
    }

    /// {代表元: 集合} のマップを返す
    ///
    /// - 時間計算量: $`O(N)`$
    pub fn enum_groups(&mut self) -> HashMap<usize, Vec<usize>> {
        (0..self.n).fold(HashMap::default(), |mut map, i| {
            let root = self.root(i);
            map.entry(root).or_default().push(i);
            map
        })
    }
}

impl UnionFindMonoid<()> {
    /// 新しいUnionFind木を生成する
    pub fn new(n: usize) -> Self {
        UnionFindMonoid {
            n,
            parent: vec![NEG1; n],
            value: vec![None; n],
            count: n,
        }
    }
}

impl<M: CommutativeMonoid> From<Vec<M::Val>> for UnionFindMonoid<M> {
    fn from(value: Vec<M::Val>) -> Self {
        let N = value.len();
        UnionFindMonoid {
            n: N,
            parent: vec![NEG1; N],
            value: value.into_iter().map(Some).collect(),
            count: N,
        }
    }
}

impl<M: CommutativeMonoid> FromIterator<M::Val> for UnionFindMonoid<M> {
    fn from_iter<T: IntoIterator<Item = M::Val>>(iter: T) -> Self {
        UnionFindMonoid::from(iter.into_iter().collect::<Vec<_>>())
    }
}

impl<M> Debug for UnionFindMonoid<M>
where
    M: CommutativeMonoid,
    M::Val: Debug + Clone,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let mut uf = UnionFindMonoid::<M> {
            n: self.n,
            parent: self.parent.clone(),
            value: self.value.clone(),
            count: self.count,
        };
        let groups = uf.enum_groups();

        f.debug_map().entries(groups).finish()
    }
}