competitive_library/structure/
weighted_disjoint_set_union.rs

1//! Weighted Union Find
2use std::collections::{HashMap, HashSet};
3#[derive(Debug, Clone)]
4enum Node {
5    Root(usize),
6    Child((usize, i64)),
7}
8///UnionFind
9#[derive(Clone, Debug)]
10pub struct WeightedDisjointSetUnion {
11    uf: Vec<Node>,
12}
13
14impl WeightedDisjointSetUnion {
15    pub fn new(n: usize) -> WeightedDisjointSetUnion {
16        WeightedDisjointSetUnion {
17            uf: vec![Node::Root(1); n],
18        }
19    }
20
21    pub fn root(&mut self, target: usize) -> (usize, i64) {
22        match self.uf[target] {
23            Node::Root(_) => (target, 0),
24            Node::Child((par, w)) => {
25                let (root, root_w) = self.root(par);
26                self.uf[target] = Node::Child((root, root_w + w));
27                (root, root_w + w)
28            }
29        }
30    }
31    pub fn unite(&mut self, x: usize, y: usize, weight: i64) -> bool {
32        let rx = self.root(x);
33        let ry = self.root(y);
34        if rx.0 == ry.0 {
35            return false;
36        }
37        let size_x = self.size(x);
38        let size_y = self.size(y);
39        let w = weight + ry.1 - rx.1;
40
41        let (rx, ry, w) = if size_x > size_y {
42            (rx, ry, -w)
43        } else {
44            (ry, rx, w)
45        };
46
47        self.uf[rx.0] = Node::Root(size_x + size_y);
48        self.uf[ry.0] = Node::Child((rx.0, w));
49
50        true
51    }
52    pub fn is_same(&mut self, x: usize, y: usize) -> bool {
53        self.root(x).0 == self.root(y).0
54    }
55    pub fn size(&mut self, x: usize) -> usize {
56        let (root, _) = self.root(x);
57        match self.uf[root] {
58            Node::Root(size) => size,
59            Node::Child(_) => 0,
60        }
61    }
62    pub fn get_diff(&mut self, x: usize, y: usize) -> i64 {
63        let rx = self.root(x);
64        let ry = self.root(y);
65        assert_eq!(rx.0, ry.0);
66        rx.1 - ry.1
67    }
68    pub fn get_same_group(&mut self, x: usize) -> HashSet<usize> {
69        let root = self.root(x);
70        let mut g = HashSet::new();
71        for i in 0..self.uf.len() {
72            if root.0 == self.root(i).0 {
73                g.insert(i);
74            }
75        }
76        g
77    }
78    pub fn get_all_groups(&mut self) -> HashMap<usize, HashSet<(usize, i64)>> {
79        let mut map: HashMap<usize, HashSet<(usize, i64)>> = HashMap::new();
80        for i in 0..self.uf.len() {
81            let (root, w) = self.root(i);
82
83            map.entry(root).or_default().insert((i, w));
84        }
85        map
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use super::*;
92
93    #[test]
94    fn test_dsu() {
95        let mut d = WeightedDisjointSetUnion::new(5);
96        d.unite(0, 1, 1);
97        assert!(d.is_same(0, 1));
98        d.unite(1, 2, 1);
99        assert!(d.is_same(0, 2));
100        assert_eq!(d.size(0), 3);
101        assert!(!d.is_same(0, 3));
102
103        d.unite(0, 3, 3);
104        d.unite(1, 4, 3);
105
106        for i in 0..5 {
107            assert_eq!(d.get_diff(0, i), i as i64);
108        }
109    }
110    #[test]
111    fn test_dsu_w2() {
112        let mut dsu = WeightedDisjointSetUnion::new(16);
113        dsu.unite(0, 1, 2);
114        dsu.unite(2, 3, 4);
115        dsu.unite(4, 5, 5);
116        dsu.unite(6, 7, 5);
117        dsu.unite(8, 9, 8);
118        dsu.unite(10, 11, 5);
119        dsu.unite(12, 13, 8);
120        dsu.unite(14, 15, 7);
121
122        dsu.unite(15, 13, -16);
123        dsu.unite(11, 9, -9);
124        dsu.unite(7, 5, -53);
125        dsu.unite(3, 1, -9);
126
127        dsu.unite(0, 15, 973);
128        dsu.unite(4, 11, 81);
129
130        dsu.unite(3, 4, 853);
131
132        let ans = [
133            0, 2, 7, 11, 864, 869, 917, 922, 928, 936, 940, 945, 949, 957, 966, 973,
134        ];
135        for (i, &v) in ans.iter().enumerate() {
136            assert_eq!(dsu.get_diff(0, i), v);
137        }
138    }
139}
competitive_library/structure/weighted_disjoint_set_union.rs

competitive_library/structure/
weighted_disjoint_set_union.rs
