Optimize water distribution in a village

Union-Find Graph Theory Heap (Priority Queue) Minimum Spanning Tree

65.5% acceptance

Mar 31, 2026

1207

There are n houses in a village. We want to supply water for all the houses by building wells and laying pipes. For each house i, we can either build a well inside it directly with cost wells[i - 1] (note the -1 due to 0-indexing), or pipe in water from another well to it. The costs to lay pipes between houses are given by the array pipes where each pipes[j] = [house1j, house2j, costj] represents the cost to connect house1j and house2j together using a pipe. Connections are bidirectional, and there could be multiple valid connections between the same two houses with different costs. Return the minimum total cost to supply water to all houses.

Solution

Rust

Time O(n)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn min_cost_to_supply_water(n: i32, wells: Vec<i32>, pipes: Vec<Vec<i32>>) -> i32 {
    // Add virtual node 0 connected to each house i with cost wells[i-1]
    let n = n as usize;
    let mut edges: Vec<(i32, usize, usize)> = Vec::new();
    for i in 0..n {
      edges.push((wells[i], 0, i + 1));
    }
    for p in &pipes {
      edges.push((p[2], p[0] as usize, p[1] as usize));
    }
    edges.sort_unstable_by_key(|e| e.0);

    let mut parent: Vec<usize> = (0..=n).collect();
    let mut rank = vec![0u8; n + 1];

    fn find(parent: &mut Vec<usize>, x: usize) -> usize {
      if parent[x] != x {
        parent[x] = find(parent, parent[x]);
      }
      parent[x]
    }

    let mut cost = 0;
    let mut edges_used = 0;
    for (w, a, b) in &edges {
      let ra = find(&mut parent, *a);
      let rb = find(&mut parent, *b);
      if ra != rb {
        if rank[ra] < rank[rb] {
          parent[ra] = rb;
        } else if rank[ra] > rank[rb] {
          parent[rb] = ra;
        } else {
          parent[rb] = ra;
          rank[ra] += 1;
        }
        cost += w;
        edges_used += 1;
        if edges_used == n {
          break;
        }
      }
    }
    cost
  }
}