Refactor the code by introducing the UnionFind class

Author: rayworks

src/main/java/org/sean/array/IslandCounter.java

@@ -1,84 +1,23 @@
 package org.sean.array;
 
+import org.sean.graph.UnionFind;
+
 /***
  * 200. Number of Islands
  */
 public class IslandCounter {
-    private int[] parent;   // parent[i] = parent of i
-    private int[] size;     // size[i] = number of sites in subtree rooted at i
-    private int count;      // number of components
-
-    private void init(int n) {
-        count = n;
-        parent = new int[n];
-        size = new int[n];
-        for (int i = 0; i < n; i++) {
-            parent[i] = i;
-            size[i] = 1;
-        }
-    }
-
-    /**
-     * The number of components.
-     */
-    private int count() {
-        return count;
-    }
-
-    private int find(int p) {
-        validate(p);
-        while (p != parent[p])
-            p = parent[p];
-        return p;
-    }
-
-    // validate that p is a valid index
-    private void validate(int p) {
-        int n = parent.length;
-        if (p < 0 || p >= n) {
-            throw new IllegalArgumentException("index " + p + " is not between 0 and " + (n - 1));
-        }
-    }
-
-    private boolean connected(int p, int q) {
-        return find(p) == find(q);
-    }
-
-    private void union(int p, int q) {
-        int rootP = find(p);
-        int rootQ = find(q);
-        if (rootP == rootQ) return;
-
-        // make smaller root point to larger one
-        if (size[rootP] < size[rootQ]) {
-            parent[rootP] = rootQ;
-            size[rootQ] += size[rootP];
-        } else {
-            parent[rootQ] = rootP;
-            size[rootP] += size[rootQ];
-        }
-        count--;
-    }
-
-    private void link(int p, int q) {
-        if (!connected(p, q)) {
-            union(p, q);
-        }
-    }
+    private UnionFind unionFind;
 
     public int numIslands(char[][] grid) {
-        if (grid == null || grid.length == 0)
-            return 0;
+        if (grid == null || grid.length == 0) return 0;
         int row = grid.length;
         int col = grid[0].length;
 
         int n = row * col;
         int p = -1;
         int q = -1;
 
-        // take advantage of WeightedQuickUnionUF from Algs4
-        // https://algs4.cs.princeton.edu/15uf/WeightedQuickUnionUF.java.html
-        init(n);
+        unionFind = new UnionFind(n);
 
         int zeroCnt = 0;
 
@@ -90,15 +29,15 @@ public int numIslands(char[][] grid) {
                     if (i + 1 < row) {
                         if (grid[i + 1][j] == '1') {
                             int p2 = (i + 1) * col + j;
-                            link(pos, p2);
+                            unionFind.union(pos, p2);
 
                             System.out.println(String.format(">>> Evaluate pos [%d, %d]", pos, p2));
                         }
                     }
                     if (j + 1 < col) {
                         if (grid[i][j + 1] == '1') {
                             int p2 = i * col + j + 1;
-                            link(pos, p2);
+                            unionFind.union(pos, p2);
 
                             System.out.println(String.format(">>> Evaluate pos [%d, %d]", pos, p2));
                         }
@@ -109,6 +48,6 @@ public int numIslands(char[][] grid) {
             }
         }
 
-        return count() - zeroCnt;
+        return unionFind.count() - zeroCnt;
     }
 }

src/main/java/org/sean/graph/UnionFind.java

@@ -0,0 +1,78 @@
+package org.sean.graph;
+
+import java.util.Arrays;
+import java.util.LinkedList;
+import java.util.List;
+
+// Case Study: Union-Find
+// https://algs4.cs.princeton.edu/15uf/
+public class UnionFind {
+    private int components;
+    private final int[] parents;
+    private final int[] sz;
+
+    public UnionFind(int n) {
+        components = n;
+
+        parents = new int[n];
+        for (int i = 0; i < n; i++) {
+            parents[i] = i; // self
+        }
+
+        sz = new int[n];
+        Arrays.fill(sz, 1);
+    }
+
+    private int topParent(int p) {
+        int curr = p;
+        while (parents[curr] != curr) {
+            curr = parents[curr];
+        }
+        return curr;
+    }
+
+    public void union(int p, int q) {
+        int pt1 = topParent(p);
+        int pt2 = topParent(q);
+        if (pt1 != pt2) {
+            // weighted quick union:
+            // connect the smaller tree to the larger
+            if (sz[pt1] > sz[pt2]) {
+                parents[pt2] = pt1;
+                sz[pt1] += sz[pt2];
+            } else {
+                parents[pt1] = pt2;
+                sz[pt2] += sz[pt1];
+            }
+
+            --components;
+        }
+    }
+
+    public boolean isConnected(int p, int q) {
+        return topParent(p) == topParent(q);
+    }
+
+    public int count() {
+        return components;
+    }
+
+    public static void main(String[] args) {
+        List<int[]> arr = new LinkedList<>();
+        arr.add(new int[] {4, 3});
+        arr.add(new int[] {3, 8});
+        arr.add(new int[] {6, 5});
+        arr.add(new int[] {9, 4});
+        arr.add(new int[] {2, 1});
+        arr.add(new int[] {5, 0});
+        arr.add(new int[] {7, 2});
+        arr.add(new int[] {6, 1});
+
+        UnionFind unionFind = new UnionFind(10);
+        for (int[] pair : arr) {
+            unionFind.union(pair[0], pair[1]);
+        }
+
+        System.out.println("Components : " + unionFind.count());
+    }
+}