Optimize Water Distribution in a Village
Problem
There are n houses in a village. We want to supply water to every house, either by building a well directly in house i for cost wells[i], or by laying a pipe between two houses where pipes[j] = [house1, house2, cost] connects them for that cost (a connected house can pass water to others). Return the minimum total cost to supply water to all n houses.
n = 3, wells = [1, 2, 2], pipes = [[1,2,1], [2,3,1]]3def min_cost_to_supply_water(n, wells, pipes):
edges = [(c, 0, i + 1) for i, c in enumerate(wells)]
for u, v, c in pipes:
edges.append((c, u, v))
edges.sort()
parent = list(range(n + 1))
def find(x):
while parent[x] != x:
parent[x] = parent[parent[x]]
x = parent[x]
return x
total, used = 0, 0
for c, u, v in edges:
ru, rv = find(u), find(v)
if ru != rv:
parent[ru] = rv
total += c
used += 1
if used == n:
break
return total
function minCostToSupplyWater(n, wells, pipes) {
const edges = wells.map((c, i) => [c, 0, i + 1]);
for (const [u, v, c] of pipes) edges.push([c, u, v]);
edges.sort((a, b) => a[0] - b[0]);
const parent = Array.from({ length: n + 1 }, (_, i) => i);
function find(x) {
while (parent[x] !== x) {
parent[x] = parent[parent[x]];
x = parent[x];
}
return x;
}
let total = 0, used = 0;
for (const [c, u, v] of edges) {
const ru = find(u), rv = find(v);
if (ru !== rv) {
parent[ru] = rv;
total += c;
if (++used === n) break;
}
}
return total;
}
class Solution {
int[] parent;
public int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
List<int[]> edges = new ArrayList<>();
for (int i = 0; i < n; i++) edges.add(new int[]{ wells[i], 0, i + 1 });
for (int[] p : pipes) edges.add(new int[]{ p[2], p[0], p[1] });
edges.sort((a, b) -> a[0] - b[0]);
parent = new int[n + 1];
for (int i = 0; i <= n; i++) parent[i] = i;
int total = 0, used = 0;
for (int[] e : edges) {
int ru = find(e[1]), rv = find(e[2]);
if (ru != rv) {
parent[ru] = rv;
total += e[0];
if (++used == n) break;
}
}
return total;
}
int find(int x) {
while (parent[x] != x) { parent[x] = parent[parent[x]]; x = parent[x]; }
return x;
}
}
vector<int> parent;
int find(int x) {
while (parent[x] != x) { parent[x] = parent[parent[x]]; x = parent[x]; }
return x;
}
int minCostToSupplyWater(int n, vector<int>& wells, vector<vector<int>>& pipes) {
vector<array<int, 3>> edges;
for (int i = 0; i < n; i++) edges.push_back({ wells[i], 0, i + 1 });
for (auto& p : pipes) edges.push_back({ p[2], p[0], p[1] });
sort(edges.begin(), edges.end());
parent.resize(n + 1);
for (int i = 0; i <= n; i++) parent[i] = i;
int total = 0, used = 0;
for (auto& e : edges) {
int ru = find(e[1]), rv = find(e[2]);
if (ru != rv) {
parent[ru] = rv;
total += e[0];
if (++used == n) break;
}
}
return total;
}
Explanation
The clever trick is to invent a virtual water source, call it node 0. Building a well in house i becomes nothing more than an edge from this source to house i with cost wells[i]. Now "supply water to every house" simply means "connect all houses (and the source) into one tree" — a classic minimum spanning tree problem.
We gather all edges: one per well from node 0, plus all the pipes. Then we run Kruskal's algorithm — sort every edge by cost and greedily add the cheapest edge that joins two pieces that are not yet connected.
To know whether an edge connects two separate pieces, we use union-find. For an edge, if find(u) != find(v) the endpoints are in different groups, so we take it, union them, and add its cost. If they already share a root, the edge would form a cycle and we skip it.
We stop once we have added n useful edges, which is exactly enough to tie all n houses to the source.
Example: n = 3, wells = [1,2,2], pipes = [[1,2,1],[2,3,1]]. The cheapest edges are well-at-house-1 (cost 1), pipe 1–2 (cost 1), pipe 2–3 (cost 1). Adding those three connects everything for a total of 3.