Water Every House
Problem
A village has n houses numbered 1 to n. House i can get water from a well of its own at cost wells[i - 1], or through pipes: a pipe [a, b, c] joins houses a and b in both directions at cost c. Return the smallest total cost that gives every house water, either from its own well or through a chain of pipes leading to a house that has a well.
Examples
Input: wells = [3, 4, 2], pipes = [[1, 2, 1], [2, 3, 5]]
Output: 6
Why: wells at houses 1 and 3, plus the pipe from 1 to 2Input: wells = [1, 1], pipes = [[1, 2, 10]]
Output: 2
Why: two cheap wells beat the expensive pipeInput: wells = [5], pipes = []
Output: 5
Why: edge case, one house and no pipesHints
0 / 3
If every house had to be joined by pipes, this would be a minimum spanning tree question. The wells are what make it different: they are a second way to connect a house.
Imagine an extra node standing for the ground water, joined to every house by an edge whose cost is that house's well. A house with a well is then just a house connected to that node.
Add the extra node, then build a minimum spanning tree over all n + 1 nodes with Prim's algorithm: start at the extra node, keep candidate edges in a min-heap, and absorb the cheapest edge that reaches a new node until every node is in.
Solution
With a node 0 for the ground water, a well at house i is the edge from 0 to i, and a plan that waters every house is exactly a set of edges that connects all n + 1 nodes. The cheapest connecting set is a minimum spanning tree, because a cycle could always drop its most expensive edge. Prim's algorithm grows the tree from node 0 with a min-heap of candidate edges and skips stale entries whose far end is already inside. Time is O((n + m) log(n + m)) for m pipes, and space is O(n + m).
import heapq
def water_cost(wells, pipes):
adj = [[] for _ in range(len(wells) + 1)] # node 0 is the ground water
for house, cost in enumerate(wells, 1):
adj[0].append((cost, house)) # a well is an edge to node 0
for a, b, cost in pipes:
adj[a].append((cost, b))
adj[b].append((cost, a))
inside, heap, total = set(), [(0, 0)], 0
while heap:
cost, node = heapq.heappop(heap) # the cheapest edge out of the tree
if node in inside:
continue # stale: both ends already inside
inside.add(node)
total += cost
for edge in adj[node]:
if edge[1] not in inside:
heapq.heappush(heap, edge)
return total
print(water_cost([3, 4, 2], [[1, 2, 1], [2, 3, 5]])) # -> 6
print(water_cost([1, 1], [[1, 2, 10]])) # -> 2
print(water_cost([5], [])) # -> 5Stuck on the idea rather than the code? Prim's Spanning Tree covers it.