Optimize Water Distribution in a Village

hard graph minimum spanning tree union find

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.

Inputn = 3, wells = [1, 2, 2], pipes = [[1,2,1], [2,3,1]]
Output3
Build a well at house 1 (cost 1), then lay pipes 1–2 (cost 1) and 2–3 (cost 1). Total = 3.

def 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;
}
Time: O(E log E) Space: O(n + E)