dsalgo/
minimum_spanning_tree_prim_sparse.rs

1// TODO
2// use generic minimum queue
3// use generic edge type
4pub fn mst_prim_sparse(
5    v_size: usize,
6    undirected_edges: &[(usize, usize, i64)],
7) -> Vec<usize> {
8    use std::cmp::Reverse;
9
10    let mut g = vec![vec![]; v_size];
11
12    for (i, &(u, v, w)) in undirected_edges.iter().enumerate() {
13        g[u].push((v, w, i));
14
15        g[v].push((u, w, i));
16    }
17
18    let mut hq = std::collections::BinaryHeap::new();
19
20    hq.push((Reverse(0), 0)); // weight, to
21    let mut eid_to: Vec<Option<usize>> = vec![None; v_size];
22
23    let mut visited = vec![false; v_size];
24
25    while let Some((_, u)) = hq.pop() {
26        if visited[u] {
27            continue;
28        }
29
30        visited[u] = true;
31
32        for &(v, w, i) in &g[u] {
33            if visited[v]
34                || eid_to[v].is_some()
35                    && w >= undirected_edges[eid_to[v].unwrap()].2
36            {
37                continue;
38            }
39
40            hq.push((Reverse(w), v));
41
42            eid_to[v] = Some(i);
43        }
44    }
45
46    eid_to.into_iter().filter_map(|e| e).collect()
47}
48
49// TODO
50#[cfg(test)]
51
52mod tests {
53
54    #[test]
55
56    fn test() {}
57}