Optimize Water Distribution in a Village

Tested Python solution for LeetCode 1168 with 19 pytest cases. Generate a practice environment with lcpy.

LeetCode 1168, Hard. Topics: Union Find, Graph, Minimum Spanning Tree, Heap (Priority Queue). View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 19 parametrized pytest cases, and a playground notebook:

lcpy gen -n 1168   # by problem number
lcpy gen -s optimize_water_distribution_in_a_village   # by problem name

​

Problem

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] = [house1_j, house2_j, cost_j] represents the cost to connect house1_j and house2_j 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.

​

Examples

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.

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 <= 10^4
• wells.length == n
• 0 <= wells[i] <= 10^5
• 1 <= pipes.length <= 10^4
• pipes[j].length == 3
• 1 <= house1_j, house2_j <= n
• 0 <= cost_j <= 10^5
• house1_j != house2_j

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

class Solution:
    # Time: O((m + n) log(m + n)) where m = len(pipes), n = len(wells)
    # Space: O(n + m)
    def min_cost_to_supply_water(self, n: int, wells: list[int], pipes: list[list[int]]) -> int:
        # Virtual well node 0: connecting house i to it costs wells[i - 1].
        edges = [(w, 0, i + 1) for i, w in enumerate(wells)]
        edges += [(c, a, b) for a, b, c in pipes]
        edges.sort()

        parent = list(range(n + 1))

        def find(x: int) -> int:
            while parent[x] != x:
                parent[x] = parent[parent[x]]
                x = parent[x]
            return x

        total = 0
        components = n + 1
        for cost, a, b in edges:
            ra, rb = find(a), find(b)
            if ra == rb:
                continue
            parent[ra] = rb
            total += cost
            components -= 1
            if components == 1:
                break
        return total

​

Complexity

Time	Space
O((m + n) log(m + n)) where m = len(pipes), n = len(wells)	O(n + m)

​

Tags