LeetCode Solutions Blog

Optimize Water Distribution in a Village

Number: 1144

Difficulty: Hard

Paid? Yes

Problem Description

There are n houses in a village. Each house can either have a well built directly in it with a given cost, or be connected to another house via pipes (with their own costs) to receive water. The goal is to supply water to all houses with minimum total cost. Note that a well in any house is equivalent to connecting that house to a virtual water source.

Key Insights

Transform the problem by introducing a virtual node (node 0) where connecting house i to this node costs wells[i - 1].
Build a graph where the edges are both the actual pipes between houses and edges from the virtual node to each house representing the well cost.
Use a Minimum Spanning Tree (MST) approach (either Kruskal’s or Prim’s algorithm) to connect all houses with the minimum cost.
The MST will automatically choose the cheaper option between building a well and laying a connecting pipe.

Space and Time Complexity

Time Complexity: O((n + p) log(n + p)), where n is the number of houses and p is the number of pipes (for sorting the edges in the MST algorithm). Space Complexity: O(n + p) to store the union-find data structures and the edge list.

Solution

We solve the problem by modeling it as an MST problem. First, we add a virtual node (0) to represent a water source. Then, we connect this node to each house i with an edge that has a cost equal to wells[i - 1]. Next, we add all the given pipe connections between houses to the graph. With this graph setup, the problem reduces to finding the minimum spanning tree that connects all nodes (including the virtual node). We can then use Kruskal’s algorithm with a union-find data structure to build this MST. The union-find data structure helps efficiently determine whether adding an edge will create a cycle, while sorting the edges ensures that we always add the smallest available edge that connects two components of the graph.

Code Solutions

# Python solution using Kruskal's algorithm and Union-Find

# Define a class to perform Union-Find (Disjoint Set Union)
class UnionFind:
    def __init__(self, size):
        # Initialize parent pointers and rank array for union by rank
        self.parent = list(range(size))
        self.rank = [0] * size

    def find(self, x):
        # Path compression to find root of x
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x, y):
        # Union two sets if roots are different; returns True if union performed
        rootX = self.find(x)
        rootY = self.find(y)
        if rootX == rootY:
            return False  # They are already connected
        # Attach the tree with lower rank to the root of the tree with higher rank
        if self.rank[rootX] < self.rank[rootY]:
            self.parent[rootX] = rootY
        elif self.rank[rootX] > self.rank[rootY]:
            self.parent[rootY] = rootX
        else:
            self.parent[rootY] = rootX
            self.rank[rootX] += 1
        return True

def minCostToSupplyWater(n, wells, pipes):
    # Build the edge list, starting with edges from a virtual node 0 to every house
    edges = []
    for i in range(1, n + 1):
        # (cost, house, virtual node 0)
        edges.append((wells[i - 1], i, 0))
    # Add all pipe edges to the list – note these are bidirectional
    for house1, house2, cost in pipes:
        edges.append((cost, house1, house2))
    # Sort edges based on cost
    edges.sort()
    
    uf = UnionFind(n + 1)
    total_cost = 0
    
    # Process each edge: if it connects two disjoint sets, add it to the MST
    for cost, house1, house2 in edges:
        if uf.union(house1, house2):
            total_cost += cost
    
    return total_cost

# Example usage:
n = 3
wells = [1, 2, 2]
pipes = [[1, 2, 1], [2, 3, 1]]
print(minCostToSupplyWater(n, wells, pipes))