LeetCode 1168. Optimize Water Distribution in a Village

Problem Statement

Union Find on minimum spanning tree.

impl Solution {
    fn find(x : usize, uf: &mut Vec<usize>) ->usize {
        if x == uf[x as usize] {
            return x;
        } else {
            uf[x] = Self::find(uf[x], uf);
            return uf[x];
        }
    }

    pub fn min_cost_to_supply_water(n: i32, wells: Vec<i32>, pipes: Vec<Vec<i32>>) -> i32 {
        let mut uf = vec![0; n as usize + 1];
        for i in 0..n + 1 {
            uf[i as usize] = i as usize;
        }
        let mut p = pipes;
        for i in 0..n {
            p.push(vec![0, i + 1, wells[i as usize]]);
        }
        p.sort_by_key(|x| x[2]);
        let mut ans = 0;
        for pipe in p {
            let (fx, fy) = (Self::find(pipe[0] as usize, &mut uf),
                                        Self::find(pipe[1] as usize, &mut uf));
            if fx == fy {
                continue;
            }
            ans += pipe[2];
            uf[fx] = fy;
        }
        return ans;
    }
}