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.

思路：跟Minimum Spanning Tree 很类似，这题多的就是一个wells，那么按照提示，我们可以把wells也转换成边，方法就是虚拟所有node跟虚拟node去link，然后第一个边需要加的就是wells最小的那条边；然后剩下的就是union find的每次去connect；

Add one virtual node (index n) to change wells to pipes, find minimum wells and connect ( mini_index, n)，sort all edges (pipes + wells) , using Union Find to connect using minimum cost.
Note: pipes[0], pipes[1] are not index, need to -1 so that change to node index;

class Solution {
    
    private class Connection {
        public int a;
        public int b;
        public int cost;
        public Connection(int a, int b, int cost) {
            this.a = a;
            this.b = b;
            this.cost = cost;
        }
    }
    
    private class ConnectionComparator implements Comparator<Connection> {
        @Override
        public int compare(Connection a, Connection b) {
            return a.cost - b.cost;
        }
    }
    
    private class UnionFind {
        private int[] father;
        private int count;
        
        public UnionFind(int n) {
            father = new int[n+1];
            for(int i = 0; i <= n; i++) {
                father[i] = i;
            }
            count = n;
        }
        
        public void union(int a, int b) {
            int root_a = find(a);
            int root_b = find(b);
            if(root_a != root_b) {
                father[root_a] = root_b;
                count--;
            }
        }
        
        public int find(int x) {
            int j = x;
            while(father[j] != j) {
                j = father[j];
            }
            
            // path compression;
            while(x != j) {
                int fx = father[x];
                father[x] = j;
                x = fx;
            }
            
            return j;
        }
    }
    
    public int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
        List<Connection> list = new ArrayList<Connection>();
        
        int minwell = Integer.MAX_VALUE;
        int minwell_index = -1;
        for(int i = 0;  i < wells.length; i++) {
            if(minwell > wells[i]) {
                minwell = wells[i];
                minwell_index = i;
            }
            list.add(new Connection(i, n, wells[i]));
        }
        
        for(int i = 0; i < pipes.length; i++) {
            list.add(new Connection(pipes[i][0]-1, pipes[i][1]-1, pipes[i][2]));
        }
        
        int cost = 0;
        UnionFind uf = new UnionFind(n + 1);
        uf.union(minwell_index, n);
        cost += minwell;
        
        Collections.sort(list, new ConnectionComparator());
        for(Connection connection : list) {
            int a = connection.a;
            int b = connection.b;
            if(uf.find(a) != uf.find(b)) {
                uf.union(a, b);
                cost += connection.cost;
            }
        }
        return cost;
    }
}