Merge pull request #187 from ia7ck/uf-unite-result

[uf] unite() の返り値の情報量を増やす

Author: ia7ck

libs/union_find/src/lib.rs

@@ -5,6 +5,12 @@ pub struct UnionFind {
     groups: usize,
 }
 
+#[derive(Clone, Debug)]
+pub struct UniteResult {
+    pub new_root: usize,
+    pub child: usize,
+}
+
 #[derive(Clone, Copy, Debug)]
 enum NodeKind {
     Root { size: usize },
@@ -61,47 +67,53 @@ impl UnionFind {
 
     /// 頂点 `i` の属する連結成分と頂点 `j` の属する連結成分をつなげます。
     ///
-    /// 呼び出し前に別の連結成分だった場合 true を、同じ連結成分だった場合 false を返します。
+    /// もともと `i` と `j` が同じ連結成分だった場合は `None` を返します。
     ///
     /// # Examples
     ///
     /// ```
     /// use union_find::UnionFind;
     /// let mut uf = UnionFind::new(6);
-    /// assert!(uf.unite(0, 1));
-    /// assert!(uf.unite(1, 2));
-    /// assert!(uf.unite(3, 4));
+    /// assert!(uf.unite(0, 1).is_some());
+    /// assert!(uf.unite(1, 2).is_some());
+    /// assert!(uf.unite(3, 4).is_some());
     ///
     /// // [(0, 1, 2), (3, 4), (5)]
-    /// assert!(!uf.unite(0, 2));
-    /// assert!(!uf.unite(3, 3));
+    /// assert!(uf.unite(0, 2).is_none());
+    /// assert!(uf.unite(3, 3).is_none());
     ///
-    /// assert!(uf.unite(4, 5));
+    /// assert!(uf.unite(4, 5).is_some());
     /// ```
-    pub fn unite(&mut self, i: usize, j: usize) -> bool {
+    pub fn unite(&mut self, i: usize, j: usize) -> Option<UniteResult> {
         let i = self.find(i);
         let j = self.find(j);
         if i == j {
-            return false;
+            return None;
         }
 
         match (self.nodes[i], self.nodes[j]) {
             (NodeKind::Root { size: i_size }, NodeKind::Root { size: j_size }) => {
+                self.groups -= 1;
                 let total = i_size + j_size;
                 // マージテク
                 if i_size >= j_size {
                     self.nodes[j] = NodeKind::Child { parent: i };
                     self.nodes[i] = NodeKind::Root { size: total };
+                    Some(UniteResult {
+                        new_root: i,
+                        child: j,
+                    })
                 } else {
                     self.nodes[i] = NodeKind::Child { parent: j };
                     self.nodes[j] = NodeKind::Root { size: total };
+                    Some(UniteResult {
+                        new_root: j,
+                        child: i,
+                    })
                 }
             }
             _ => unreachable!(),
         }
-
-        self.groups -= 1;
-        true
     }
 
     /// 頂点 `i` の属する連結成分のサイズ (頂点数) を返します。