competitive_library/structure/
disjoint_set_union.rs

1//! Union find
2use std::collections::{HashMap, HashSet};
3#[derive(Copy, Clone, Debug)]
4enum Node {
5    Root(usize),
6    Child(usize),
7}
8///UnionFind
9#[derive(Clone, Debug)]
10pub struct DisjointSetUnion {
11    nodes: Vec<Node>,
12}
13
14impl DisjointSetUnion {
15    pub fn new(n: usize) -> DisjointSetUnion {
16        DisjointSetUnion {
17            nodes: vec![Node::Root(1); n],
18        }
19    }
20
21    pub fn find_root(&mut self, target: usize) -> usize {
22        match unsafe { *self.nodes.get_unchecked(target) } {
23            Node::Root(_) => target,
24            Node::Child(parent) => {
25                let parent_index = self.find_root(parent);
26                self.nodes[target] = Node::Child(parent_index);
27                parent_index
28            }
29        }
30    }
31    pub fn unite(&mut self, x: usize, y: usize) -> bool {
32        let rx = self.find_root(x);
33        let ry = self.find_root(y);
34        if rx == ry {
35            return false;
36        }
37        let size_x = self.size(x);
38        let size_y = self.size(y);
39
40        let (i, j) = if size_x > size_y { (rx, ry) } else { (ry, rx) };
41        self.nodes[i] = Node::Root(size_x + size_y);
42        self.nodes[j] = Node::Child(i);
43
44        true
45    }
46    pub fn is_same(&mut self, x: usize, y: usize) -> bool {
47        self.find_root(x) == self.find_root(y)
48    }
49    pub fn size(&mut self, x: usize) -> usize {
50        let root = self.find_root(x);
51        match self.nodes[root] {
52            Node::Root(size) => size,
53            Node::Child(_) => 0,
54        }
55    }
56    pub fn get_same_group(&mut self, x: usize) -> HashSet<usize> {
57        let root = self.find_root(x);
58        let mut g = HashSet::new();
59        for i in 0..self.nodes.len() {
60            if root == self.find_root(i) {
61                g.insert(i);
62            }
63        }
64        g
65    }
66    pub fn get_all_groups(&mut self) -> HashMap<usize, HashSet<usize>> {
67        let mut map: HashMap<usize, HashSet<usize>> = HashMap::new();
68        for i in 0..self.nodes.len() {
69            map.entry(self.find_root(i)).or_default().insert(i);
70        }
71        map
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::*;
78
79    #[test]
80    fn test_dsu() {
81        let mut d = DisjointSetUnion::new(4);
82        d.unite(0, 1);
83        assert!(d.is_same(0, 1));
84        d.unite(1, 2);
85        assert!(d.is_same(0, 2));
86        assert_eq!(d.size(0), 3);
87        assert!(!d.is_same(0, 3));
88
89        // assert_eq!(d.get_all_groups(), vec![vec![0, 1, 2], vec![3]]);
90    }
91}
competitive_library/structure/disjoint_set_union.rs

competitive_library/structure/
disjoint_set_union.rs
