LeetCode 1168 - Optimize Water Distribution in a Village

The problem asks us to supply water to all houses in a village in the most cost-effective way. Each house has two options: either build a well in the house with a fixed cost, or connect the house to another house via pipes, where each pipe has its own construction cost.

leetcodehardunion-findgraph-theoryheap-(priority-queue)minimum-spanning-tree

LeetCode Problem 1168

Difficulty: 🔴 Hard
Topics: Union-Find, Graph Theory, Heap (Priority Queue), Minimum Spanning Tree

Solution

Problem Understanding

The problem asks us to supply water to all houses in a village in the most cost-effective way. Each house has two options: either build a well in the house with a fixed cost, or connect the house to another house via pipes, where each pipe has its own construction cost. The goal is to find a combination of wells and pipes that minimizes the total cost of supplying water to all houses.

The inputs are as follows: n is the number of houses, wells is an array where wells[i - 1] represents the cost to build a well in house i, and pipes is a list of triplets [house1, house2, cost] indicating the bidirectional pipe construction cost between house1 and house2. The output is a single integer representing the minimum total cost to supply water.

Constraints indicate that n and the number of pipes can go up to 10,000 and costs can be as high as 100,000. This tells us that brute-force approaches that explore all combinations of wells and pipes are infeasible. Important edge cases include scenarios where building a well is cheaper than any pipe connection, or where multiple pipes connect the same pair of houses but with different costs.

Approaches

A brute-force approach would consider every combination of building wells and connecting pipes. For each subset of houses with wells, we would then attempt all possible pipe configurations to supply the remaining houses. While this is guaranteed to find the minimum cost, it is computationally infeasible due to combinatorial explosion, especially since the number of houses and pipes can reach 10,000.

The key insight for an optimal solution comes from graph theory. We can treat the problem as finding a Minimum Spanning Tree (MST). We introduce a virtual node representing a water source and connect it to each house with an edge weight equal to the well construction cost. Then, we treat all pipes as regular edges between houses. Using either Kruskal's algorithm with Union-Find or Prim's algorithm with a priority queue, we can efficiently find the MST, which represents the minimal combination of wells and pipes to supply water to all houses.

Approach Time Complexity Space Complexity Notes
Brute Force O(2^n * m) O(n + m) Explore all well and pipe combinations, infeasible for n ~ 10^4
Optimal (MST) O((n + m) log (n + m)) O(n + m) Treat wells as edges to virtual node, use Kruskal's or Prim's algorithm

Algorithm Walkthrough

  1. Introduce a virtual node 0 to represent the water source. Connect this node to every house i with an edge weight equal to wells[i-1].
  2. Treat the original pipes array as edges between houses.
  3. Combine all edges into a single list: edges from virtual node to each house plus all pipes.
  4. Sort all edges by cost in ascending order. This ensures that the MST algorithm considers cheaper connections first.
  5. Initialize a Union-Find (Disjoint Set Union) data structure to track connected components.
  6. Iterate through the sorted edges. For each edge, if the two nodes are not already connected, add the edge to the MST and union their sets. Accumulate the cost.
  7. Continue until all houses are connected to the water supply. At this point, the accumulated cost is the minimum total cost.

Why it works: The MST guarantees the minimum total weight to connect all nodes in a graph. By treating wells as edges from a virtual source node, the algorithm naturally chooses the cheaper option between building a well and connecting with pipes.

Python Solution

from typing import List

class UnionFind:
    def __init__(self, n: int):
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x: int, y: int) -> bool:
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False
        if self.rank[root_x] < self.rank[root_y]:
            self.parent[root_x] = root_y
        elif self.rank[root_x] > self.rank[root_y]:
            self.parent[root_y] = root_x
        else:
            self.parent[root_y] = root_x
            self.rank[root_x] += 1
        return True

class Solution:
    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:
        edges = [[0, i + 1, wells[i]] for i in range(n)]
        edges.extend(pipes)
        edges.sort(key=lambda x: x[2])

        uf = UnionFind(n + 1)
        total_cost = 0

        for u, v, cost in edges:
            if uf.union(u, v):
                total_cost += cost

        return total_cost

The Python solution first constructs edges from the virtual node to every house and combines them with the given pipes. Sorting ensures we process the cheapest connections first. Union-Find manages connectivity and prevents cycles while accumulating the total minimum cost.

Go Solution

package main

import "sort"

type UnionFind struct {
	parent []int
	rank   []int
}

func NewUnionFind(n int) *UnionFind {
	parent := make([]int, n)
	rank := make([]int, n)
	for i := 0; i < n; i++ {
		parent[i] = i
	}
	return &UnionFind{parent, rank}
}

func (uf *UnionFind) Find(x int) int {
	if uf.parent[x] != x {
		uf.parent[x] = uf.Find(uf.parent[x])
	}
	return uf.parent[x]
}

func (uf *UnionFind) Union(x, y int) bool {
	rootX, rootY := uf.Find(x), uf.Find(y)
	if rootX == rootY {
		return false
	}
	if uf.rank[rootX] < uf.rank[rootY] {
		uf.parent[rootX] = rootY
	} else if uf.rank[rootX] > uf.rank[rootY] {
		uf.parent[rootY] = rootX
	} else {
		uf.parent[rootY] = rootX
		uf.rank[rootX]++
	}
	return true
}

func minCostToSupplyWater(n int, wells []int, pipes [][]int) int {
	edges := make([][]int, n)
	for i := 0; i < n; i++ {
		edges[i] = []int{0, i + 1, wells[i]}
	}
	edges = append(edges, pipes...)

	sort.Slice(edges, func(i, j int) bool {
		return edges[i][2] < edges[j][2]
	})

	uf := NewUnionFind(n + 1)
	totalCost := 0
	for _, edge := range edges {
		if uf.Union(edge[0], edge[1]) {
			totalCost += edge[2]
		}
	}
	return totalCost
}

In Go, the implementation mirrors Python, with attention to slice allocation and sort function syntax. Union-Find is implemented as a struct with methods.

Worked Examples

Example 1

Input: n = 3, wells = [1,2,2], pipes = [[1,2,1],[2,3,1]]

Edges after adding virtual node:

[[0,1,1],[0,2,2],[0,3,2],[1,2,1],[2,3,1]]

Sorted edges:

[[0,1,1],[1,2,1],[2,3,1],[0,2,2],[0,3,2]]

Step-by-step Union-Find:

Edge Union? Total Cost Connected Houses
[0,1,1] Yes 1 1
[1,2,1] Yes 2 1,2
[2,3,1] Yes 3 1,2,3
[0,2,2] No 3 already connected
[0,3,2] No 3 already connected

Output: 3

Complexity Analysis

Measure Complexity Explanation
Time O((n + m) log (n + m)) Sorting all edges dominates, union-find operations are near-constant amortized
Space O(n + m) Storage for edges and union-find parent/rank arrays

Sorting is the bottleneck, and union-find with path compression ensures efficient connectivity checks.

Test Cases

# Example test cases
assert Solution().minCostToSupplyWater(3, [1,2,2], [[1,2,1],[2,3,1]]) == 3  # Example 1
assert Solution().minCostToSupplyWater(2, [1,1], [[1,2,1],[1,2,2]]) == 2    # Example 2

# Boundary tests
assert Solution().minCostToSupplyWater(2, [0,0], [[1,2,100]]) == 0           # Zero cost wells
assert Solution