Optimize Water Distribution in a Village

There are n houses in a village. We want to supply water for all the houses by building wells and laying pipes.

For each house i, we can either build a well inside it directly with cost wells[i], or pipe in water from another well to it. The costs to lay pipes between houses are given by the array pipes, where each pipes[i] = [house1, house2, cost] represents the cost to connect house1 and house2 together using a pipe. Connections are bidirectional.

Find the minimum total cost to supply water to all houses.

 

Example 1:

Input: n = 3, wells = [1,2,2], pipes = [[1,2,1],[2,3,1]]
Output: 3
Explanation: 
The image shows the costs of connecting houses using pipes.
The best strategy is to build a well in the first house with cost 1 and connect the other houses to it with cost 2 so the total cost is 3.
 

Constraints:


Solution:

class Solution {
    static class UF {
        int[] parent;
        int[] size;
        int count;
        
        public UF(int n) {
            parent = new int[n + 1];
            size = new int[n + 1];
            for (int i = 0; i <= n; i ++) {
                parent[i] = i;
                size[i] = 1;
            }
            count = n + 1;
        }
        
        public int find(int x) {
            while (parent[x] != x) {
                parent[x] = parent[parent[x]];
                x = parent[x];
            }
            return x;
        }
        
        public boolean union(int x, int y) {
            int rx = find(x);
            int ry = find(y);
            if (rx == ry) return false;
            if (size[rx] > size[ry]) {
                parent[ry] = rx;
                size[rx] += size[ry];
            } else {
                parent[ry] = rx;
                size[ry] += size[rx];
            }
            count --;
            return true;
        }
    }
    
    public int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
        PriorityQueue<int[]> pq = new PriorityQueue<int[]>((a, b) -> a[2] - b[2]);
        for (int i = 0; i < wells.length; i ++) {
            pq.offer(new int[]{i + 1, i + 1, wells[i]});
        }
        for (int[] pipe : pipes) {
            pq.offer(pipe);
        }        
        UF houses = new UF(n);
        int cost = 0;
        while (houses.count > 1) {
            int[] next = pq.poll();
            if (next[0] == next[1] && houses.union(0, next[0])) {
                cost += next[2];
            } else if (houses.union(next[0], next[1])) {
                cost += next[2];
            }
        }
        return cost;
    }
}