Optimize Water Distribution in a Village in Python

The water distribution optimization problem is a classic application of graph theory and the Minimum Spanning Tree (MST) algorithm. Given n houses, we can either build wells directly in houses or connect them via pipes. The goal is to minimize the total cost to supply water to all houses.

The problem can be modeled as finding the MST of a graph where:

  • Each house is a node

  • Building a well is represented as connecting to a virtual node (node 0)

  • Pipes between houses are edges with their respective costs

Algorithm Approach

We use Kruskal's algorithm with Union-Find data structure:

  1. Add virtual edges from node 0 to each house with well costs

  2. Sort all edges by cost

  3. Use Union-Find to detect cycles and build MST

0 Well Source 1 2 3 cost: 1 cost: 2 cost: 2 pipe: 1 pipe: 1

Implementation

class Solution:
    def __init__(self):
        self.parent = []
    
    def find(self, a):
        """Find parent with path compression"""
        if self.parent[a] == -1:
            return a
        self.parent[a] = self.find(self.parent[a])
        return self.parent[a]
    
    def union(self, a, b):
        """Union two nodes, return True if they were already connected"""
        parent_a = self.find(a)
        parent_b = self.find(b)
        
        if parent_a == parent_b:
            return True  # Cycle detected
        
        self.parent[parent_b] = parent_a
        return False
    
    def minCostToSupplyWater(self, n, wells, pipes):
        """Find minimum cost to supply water to all houses"""
        # Initialize Union-Find structure
        self.parent = [-1 for _ in range(n + 1)]
        
        # Add virtual edges from node 0 to each house (well costs)
        for i in range(len(wells)):
            pipes.append([0, i + 1, wells[i]])
        
        # Sort all edges by cost (Kruskal's algorithm)
        pipes.sort(key=lambda edge: edge[2])
        
        total_cost = 0
        edges_used = 0
        
        # Process edges in ascending order of cost
        for source, destination, cost in pipes:
            if not self.union(source, destination):
                total_cost += cost
                edges_used += 1
                
                # MST complete when we have n edges (including virtual node)
                if edges_used == n:
                    break
        
        return total_cost

# Test the solution
solution = Solution()
result = solution.minCostToSupplyWater(3, [1, 2, 2], [[1, 2, 1], [2, 3, 1]])
print("Minimum cost to supply water:", result)
Minimum cost to supply water: 3

How It Works

For the given example with n=3, wells=[1,2,2], pipes=[[1,2,1],[2,3,1]]:

  1. Add virtual edges: [0,1,1], [0,2,2], [0,3,2]

  2. All edges sorted by cost: [0,1,1], [1,2,1], [2,3,1], [0,2,2], [0,3,2]

  3. Select edges: Build well in house 1 (cost=1), connect houses 1-2 (cost=1), connect houses 2-3 (cost=1)

  4. Total minimum cost: 3

Time and Space Complexity

Aspect Complexity Explanation
Time O(E log E) Sorting edges dominates
Space O(N) Union-Find parent array

Where E = number of pipes + number of houses, N = number of houses.

Conclusion

This solution transforms the water distribution problem into a Minimum Spanning Tree problem by introducing a virtual node for well costs. The Union-Find data structure efficiently detects cycles and builds the optimal network with minimum total cost.

Updated on: 2026-03-25T09:08:31+05:30

822 Views

Kickstart Your Career

Get certified by completing the course

Get Started
Advertisements