Added dsu

Author: hotman78

src/dsu.rs

@@ -1 +1,68 @@
+// Implement (union by size) + (path compression)
+// Reference:
+// Zvi Galil and Giuseppe F. Italiano,
+// Data structures and algorithms for disjoint set union problems
+pub struct Dsu {
+    n: usize,
+    // root node: -1 * component size
+    // otherwise: parent
+    parent_or_size: Vec<i32>,
+}
 
+impl Dsu {
+    pub fn new(n: usize) -> Self {
+        Self {
+            n: n,
+            parent_or_size: vec![-1; n],
+        }
+    }
+    pub fn merge(&mut self, a: usize, b: usize) -> usize {
+        assert!(a < self.n);
+        assert!(b < self.n);
+        let (mut x, mut y) = (self.leader(a), self.leader(b));
+        if x == y {
+            return x;
+        }
+        if -self.parent_or_size[x] < -self.parent_or_size[y] {
+            std::mem::swap(&mut x, &mut y);
+        }
+        self.parent_or_size[x] += self.parent_or_size[y];
+        self.parent_or_size[y] = x as i32;
+        return x;
+    }
+
+    pub fn same(&mut self, a: usize, b: usize) -> bool {
+        assert!(a < self.n);
+        assert!(b < self.n);
+        return self.leader(a) == self.leader(b);
+    }
+    pub fn leader(&mut self, a: usize) -> usize {
+        assert!(a < self.n);
+        if self.parent_or_size[a] < 0 {
+            return a;
+        }
+        self.parent_or_size[a] = self.leader(self.parent_or_size[a] as usize) as i32;
+        self.parent_or_size[a] as usize
+    }
+    pub fn size(&mut self, a: usize) -> usize {
+        assert!(a < self.n);
+        let x = self.leader(a);
+        -self.parent_or_size[x] as usize
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn dsu_works() {
+        let mut d = Dsu::new(4);
+        d.merge(0, 1);
+        assert_eq!(d.same(0, 1), true);
+        d.merge(1, 2);
+        assert_eq!(d.same(0, 2), true);
+        assert_eq!(d.size(0), 3);
+        assert_eq!(d.same(0, 3), false);
+    }
+}