Union-Find — Optimize Water Distribution in a Village

↗ LeetCode: Optimize Water Distribution in a Village · Hard · pattern chapter →

Asked atAmazonGoogle

Given n houses, wells[i] = cost to dig a well at house i, and pipes[i] = [a, b, cost] = cost to build a pipe between houses. Provide water to all houses at min cost.

Example 1 — n=3, wells=[1,2,2], pipes=[[1,2,1],[2,3,1]] → 3 (dig at 1: 1; pipes 1-2:1 and 2-3:1) Example 2 — n=2, wells=[1,1], pipes=[] → 2 (no pipes, dig both) Example 3 — n=5, wells=[46012,72474,64965,751,33304], pipes=[[2,1,6719],[3,2,75312],[5,3,44918]] → 131704

Constraints — 2 ≤ n ≤ 10⁴; |pipes| ≤ 10⁴. Brute force enumerating which houses to well vs pipe is 2ⁿ — impossible. Kruskal on n+1 nodes is O((n+E) log(n+E)) ≈ 3·10⁵.

Progressive Hints

Try to solve it yourself first. Reveal hints only when stuck — each hint is progressively more specific.

---

⏱️Interview timer00:00

Practice under interview pressure. Start when you begin reading; stop when you finish. Save best only when you had a full working solution.

Target for interview: Medium = 20-25 min · Hard = 30-40 min · Best time persists on this device only.

Approach 1 — Brute force: try every well subset

Intuition. Enumerate every subset of houses that gets a well. For each subset, run a Steiner-tree-like MST connecting the rest to a well through pipes. Return the min.

java

// Sketch only — exponential
int bruteMinCost(int n, int[] wells, int[][] pipes) {
    int best = Integer.MAX_VALUE;
    for (int mask = 1; mask < (1 << n); mask++) {  // must have at least one well
        int wellCost = wellsSum(mask, wells);
        int connectCost = connectRestViaPipes(mask, n, pipes);
        if (connectCost == UNREACHABLE) continue;
        best = Math.min(best, wellCost + connectCost);
    }
    return best;
}

Complexity — Time O(2ⁿ · MST); Space O(n). For n=20, 2²⁰ ≈ 10⁶ and each MST call is O(E log E) — feasible but slow. Past n=25 fails. In an interview state this then flip the perspective.

Approach 2 — Add virtual node + MST (canonical)

Insight. Model wells as pipes from a virtual node 0 to each house with weight wells[i]. Now the problem is exactly a standard MST over n+1 nodes and (n + E) edges. The MST will pick the cheapest well-or-pipe combination automatically.

java

int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
    List<int[]> edges = new ArrayList<>();
    for (int i = 0; i < n; i++) edges.add(new int[]{0, i + 1, wells[i]});
    for (int[] p : pipes) edges.add(p);
    edges.sort((a, b) -> a[2] - b[2]);
    int[] parent = new int[n + 1];
    for (int i = 0; i <= n; i++) parent[i] = i;
    int cost = 0, cnt = 0;
    for (int[] e : edges) {
        int ra = find(parent, e[0]), rb = find(parent, e[1]);
        if (ra == rb) continue;
        parent[ra] = rb; cost += e[2];
        if (++cnt == n) break;
    }
    return cost;
}
int find(int[] p, int x) { return p[x] == x ? x : (p[x] = find(p, p[x])); }

Execution Trace

Virtual root — n=3

Step 1 / 3dig at 1

Complexity — Time O((n + E) log(n + E)); Space O(n + E). Say aloud in an interview: "the virtual-node trick — turning 'choose one from each group' into a standard MST — also solves Cheapest Flights With K Stops, Min Cost to Connect Islands, and Optimize Water Distribution."

Try it yourself

Complexity summary

Approach	Time	Space	Grade
Brute subset + MST	O(2ⁿ · E log E)	O(n)	Reference; TLE past n=25
Virtual node + Kruskal	O((n+E) log(n+E))	O(n+E)	Canonical

When to use which

"Free option per node" → virtual root MST.
"Free option per subset" → generalizes with more virtual nodes.
Dense pipes → Prim + heap.

✨AI Study CompanionBeta

Was this page helpful?

Union-Find — Optimize Water Distribution in a Village ​

Approach 1 — Brute force: try every well subset ​

Approach 2 — Add virtual node + MST (canonical) ​

Try it yourself ​

Complexity summary ​

When to use which ​

Related problems ​