Optimize Water Distribution in a Village

Problem#

There are n houses. House i can have its own well (cost wells[i-1]) or share water via pipes (each pipe [a, b, cost]). Return the minimum total cost to bring water to every house.

Examples#

n = 3, wells = [1,2,2], pipes = [[1,2,1],[2,3,1]] → 3 (well at 1 cost 1, pipes 1-2 and 2-3 cost 1 each).
n = 2, wells = [1,1], pipes = [[1,2,1],[1,2,8]] → 2.

Constraints#

2 <= n <= 10^4, 0 <= pipes.length <= 10^4.

Hints#

Hint

Model each well as a pipe from a virtual node 0 to the house with cost wells[i]. Then MST over n + 1 nodes.

Approach#

Add virtual node 0. For each house i create edge (0, i, wells[i-1]). Combine with the given pipes. Kruskal’s MST over n + 1 nodes — sum of accepted edge costs is the answer.

Solution#

class Solution {
    vector<int> parent, rnk;
    int find(int x) { while (parent[x] != x) { parent[x] = parent[parent[x]]; x = parent[x]; } return x; }
    bool unite(int a, int b) {
        a = find(a); b = find(b); if (a == b) return false;
        if (rnk[a] < rnk[b]) swap(a, b);
        parent[b] = a; if (rnk[a] == rnk[b]) ++rnk[a];
        return true;
    }
public:
    int minCostToSupplyWater(int n, vector<int>& wells, vector<vector<int>>& pipes) {
        vector<tuple<int,int,int>> edges;
        for (int i = 0; i < n; ++i) edges.emplace_back(wells[i], 0, i + 1);
        for (auto& p : pipes) edges.emplace_back(p[2], p[0], p[1]);
        sort(edges.begin(), edges.end());
        parent.resize(n + 1); rnk.assign(n + 1, 0);
        iota(parent.begin(), parent.end(), 0);
        int total = 0, used = 0;
        for (auto& [w, a, b] : edges) {
            if (unite(a, b)) { total += w; if (++used == n) break; }
        }
        return total;
    }
};

func minCostToSupplyWater(n int, wells []int, pipes [][]int) int {
    type edge struct{ w, a, b int }
    edges := make([]edge, 0, n+len(pipes))
    for i := 0; i < n; i++ {
        edges = append(edges, edge{wells[i], 0, i + 1})
    }
    for _, p := range pipes {
        edges = append(edges, edge{p[2], p[0], p[1]})
    }
    sort.Slice(edges, func(i, j int) bool { return edges[i].w < edges[j].w })
    parent := make([]int, n+1)
    rnk := make([]int, n+1)
    for i := range parent {
        parent[i] = i
    }
    var find func(x int) int
    find = func(x int) int {
        for parent[x] != x {
            parent[x] = parent[parent[x]]
            x = parent[x]
        }
        return x
    }
    unite := func(a, b int) bool {
        a, b = find(a), find(b)
        if a == b {
            return false
        }
        if rnk[a] < rnk[b] {
            a, b = b, a
        }
        parent[b] = a
        if rnk[a] == rnk[b] {
            rnk[a]++
        }
        return true
    }
    total, used := 0, 0
    for _, e := range edges {
        if unite(e.a, e.b) {
            total += e.w
            used++
            if used == n {
                break
            }
        }
    }
    return total
}

from typing import List

class Solution:
    def minCostToSupplyWater(self, n: int, wells: List[int], pipes: List[List[int]]) -> int:
        edges: List[tuple[int, int, int]] = []
        for i in range(n):
            edges.append((wells[i], 0, i + 1))
        for a, b, w in pipes:
            edges.append((w, a, b))
        edges.sort()
        parent = list(range(n + 1))
        rnk = [0] * (n + 1)

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

        def unite(a: int, b: int) -> bool:
            a, b = find(a), find(b)
            if a == b:
                return False
            if rnk[a] < rnk[b]:
                a, b = b, a
            parent[b] = a
            if rnk[a] == rnk[b]:
                rnk[a] += 1
            return True

        total = 0
        used = 0
        for w, a, b in edges:
            if unite(a, b):
                total += w
                used += 1
                if used == n:
                    break
        return total

function minCostToSupplyWater(n, wells, pipes) {
    const edges = [];
    for (let i = 0; i < n; i++) edges.push([wells[i], 0, i + 1]);
    for (const [a, b, w] of pipes) edges.push([w, a, b]);
    edges.sort((x, y) => x[0] - y[0]);
    const parent = Array.from({ length: n + 1 }, (_, i) => i);
    const rnk = new Array(n + 1).fill(0);
    const find = (x) => {
        while (parent[x] !== x) {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    };
    const unite = (a, b) => {
        a = find(a); b = find(b);
        if (a === b) return false;
        if (rnk[a] < rnk[b]) [a, b] = [b, a];
        parent[b] = a;
        if (rnk[a] === rnk[b]) rnk[a]++;
        return true;
    };
    let total = 0, used = 0;
    for (const [w, a, b] of edges) {
        if (unite(a, b)) {
            total += w;
            if (++used === n) break;
        }
    }
    return total;
}

class Solution {
    private int[] parent;
    private int[] rnk;

    public int minCostToSupplyWater(int n, int[] wells, int[][] pipes) {
        int[][] edges = new int[n + pipes.length][3];
        int idx = 0;
        for (int i = 0; i < n; i++) {
            edges[idx++] = new int[]{wells[i], 0, i + 1};
        }
        for (int[] p : pipes) {
            edges[idx++] = new int[]{p[2], p[0], p[1]};
        }
        Arrays.sort(edges, (a, b) -> a[0] - b[0]);
        parent = new int[n + 1];
        rnk = new int[n + 1];
        for (int i = 0; i <= n; i++) parent[i] = i;
        int total = 0, used = 0;
        for (int[] e : edges) {
            if (unite(e[1], e[2])) {
                total += e[0];
                if (++used == n) break;
            }
        }
        return total;
    }

    private int find(int x) {
        while (parent[x] != x) {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    }

    private boolean unite(int a, int b) {
        a = find(a);
        b = find(b);
        if (a == b) return false;
        if (rnk[a] < rnk[b]) {
            int t = a; a = b; b = t;
        }
        parent[b] = a;
        if (rnk[a] == rnk[b]) rnk[a]++;
        return true;
    }
}

function minCostToSupplyWater(n: number, wells: number[], pipes: number[][]): number {
    const edges: number[][] = [];
    for (let i = 0; i < n; i++) edges.push([wells[i], 0, i + 1]);
    for (const [a, b, w] of pipes) edges.push([w, a, b]);
    edges.sort((x, y) => x[0] - y[0]);
    const parent: number[] = Array.from({ length: n + 1 }, (_, i) => i);
    const rnk: number[] = new Array(n + 1).fill(0);
    const find = (x: number): number => {
        while (parent[x] !== x) {
            parent[x] = parent[parent[x]];
            x = parent[x];
        }
        return x;
    };
    const unite = (a: number, b: number): boolean => {
        a = find(a); b = find(b);
        if (a === b) return false;
        if (rnk[a] < rnk[b]) [a, b] = [b, a];
        parent[b] = a;
        if (rnk[a] === rnk[b]) rnk[a]++;
        return true;
    };
    let total = 0, used = 0;
    for (const [w, a, b] of edges) {
        if (unite(a, b)) {
            total += w;
            if (++used === n) break;
        }
    }
    return total;
}

Editorial#

The virtual-node trick converts the asymmetry between “build well” and “lay pipe” into a uniform MST problem. With n + 1 nodes (houses plus the virtual source), MST has exactly n edges — every house either connects to the source (well) or to another house (pipe). Kruskal’s runs in O(E log E).

Complexity#

Time: O((N + P) log (N + P)).
Space: O(N + P).

Problem#

Examples#

Constraints#

Hints#

Approach#

Solution#

Editorial#

Complexity#

Concept revision#

Keyboard shortcuts

Problem#

Examples#

Constraints#

Hints#

Approach#

Solution#

Editorial#

Complexity#

Concept revision#

Related#