1135. Connecting Cities With Minimum Cost

class Solution {
    class UnionFind {
        private int[] father;
        
        public UnionFind(int size) {
            father = new int[size];
            
            // Important here to initialize father
            for (int i = 0; i < size; i++) {
                father[i] = i;
            }
        }

        public int find(int x) {
            if (father[x] == x) {
                return x;
            }
            return father[x] = find(father[x]);
        }

        public void union(int a, int b) {
            int rootA = find(a);
            int rootB = find(b);
            if (rootA != rootB) {
                father[rootA] = rootB;
            }
        }

        public boolean isUnion(int a, int b) {
            int rootA = find(a);
            int rootB = find(b);
            return rootA == rootB;
        }        
    }
    
    public int minimumCost(int N, int[][] connections) {
        UnionFind uf = new UnionFind(N + 1);
        
        // 1. Sort connections by cost
        Arrays.sort(connections, new Comparator<int[]>() {
            @Override
            public int compare(int[] a, int[] b) {
                return a[2] - b[2];
            }
        });
        
        int numOfConnections = 0;
        int minimumCost = 0;
        
        // 2. Iterate connections, if the two are not connected, union connection and add cost to res
        for (int[] connection : connections) {
            if (!uf.isUnion(connection[0], connection[1])) {
                uf.union(connection[0], connection[1]);
                minimumCost += connection[2];
                numOfConnections ++;
                if (numOfConnections == N - 1) {
                    return minimumCost;
                }
            }
        }
        return -1;
    }
}