1168. 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 - 1] (note the -1 due to 0-indexing), 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[j] = [house1j, house2j, costj] represents the cost to connect house1j and house2j together using a pipe. Connections are bidirectional, and there could be multiple valid connections between the same two houses with different costs.

Return 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.

Example 2:

Input: n = 2, wells = [1,1], pipes = [[1,2,1],[1,2,2]]
Output: 2
Explanation: We can supply water with cost two using one of the three options:
Option 1:
  - Build a well inside house 1 with cost 1.
  - Build a well inside house 2 with cost 1.
The total cost will be 2.
Option 2:
  - Build a well inside house 1 with cost 1.
  - Connect house 2 with house 1 with cost 1.
The total cost will be 2.
Option 3:
  - Build a well inside house 2 with cost 1.
  - Connect house 1 with house 2 with cost 1.
The total cost will be 2.
Note that we can connect houses 1 and 2 with cost 1 or with cost 2 but we will always choose the cheapest option. 

Constraints:

• 2 <= n <= 104
• wells.length == n
• 0 <= wells[i] <= 105
• 1 <= pipes.length <= 104
• pipes[j].length == 3
• 1 <= house1j, house2j <= n
• 0 <= costj <= 105
• house1j != house2j

Solution:

class Solution {
    public int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
        // List to store all edge
        List<int[]> edges = new ArrayList<>();



        // Add well costs as edges from virtal ndoe 0 to each house
        for (int i = 0; i < wells.length; i++){
            edges.add(new int[]{0, i + 1, wells[i]});
            // [0, 1, 1]
            // [0, 2, 2]
            // [0, 3, 2]
        }


        // Add pipes as edge between houses
        for (int[] pipe : pipes){
            edges.add(pipe);
        }

        // Sort all edges by cost in ascending order
        edges.sort((a,b) -> (a[2] - b[2]));

        // Initialize Union-Find with n + 1 ndoes(including the virtual node 0)
        UnionFind uf = new UnionFind(n + 1);

        int totalCost = 0;

        // Process each edge in ascending order of cost
        for (int[] edge : edges){
            int house1 = edge[0];
            int house2 = edge[1];
            int cost = edge[2];

            // Only add the edge if it connects two unconnected componets
            if (uf.isC(house1, house2) == false){
                uf.union(house1, house2);
                totalCost += cost;
            }
        }

        return totalCost;


    }
}

class UnionFind{
    int[] root;
    int[] rank;

    public UnionFind(int n){
        root = new int[n];
        rank = new int[n];
        for (int i = 0; i < n; i++){
            root[i] = i;
            rank[i] = 1;
        }



    }

    public int find(int x){
        while(x != root[x]){
            x = root[x];
        }

        return x;
    }


    public void union(int x, int y){
        int rootX = find(x);
        int rootY = find(y);


        if (rank[rootX] < rank[rootY]){
            root[rootX] = rootY;

        }else if (rank[rootX] > rank[rootY]){
            root[rootY] = rootX;

        }else{
            root[rootY] = rootX;
            rank[rootX]++;

        }
    }

    public boolean isC(int x, int y){
        return find(x) == find(y);
    }
}

//TC: O((n+m)log(n+m))      // The number of edges is E = n + m
//SC: O(n+m)