competitive_library/graph/
minimum_spanning_tree_prim.rs

1//! Prim's minimum spanning tree algorithms.
2
3/// O(|N|^2)
4pub fn prim(g: &[Vec<i64>]) -> i64 {
5    let n = g.len();
6    let mut min_cost = vec![i64::MAX; n];
7    let mut used = vec![false; n];
8    let mut sum = 0;
9    min_cost[0] = 0;
10    loop {
11        let mut v = None;
12        for u in 0..n {
13            if used[u] || v.filter(|&x| min_cost[u] > min_cost[x]).is_some() {
14                continue;
15            }
16            v = Some(u);
17        }
18        if v.is_none() {
19            break;
20        }
21        let v = v.unwrap();
22        used[v] = true;
23        sum += min_cost[v];
24        (0..n).filter(|&u| g[v][u] != -1).for_each(|u| {
25            min_cost[u] = min_cost[u].min(g[v][u]);
26        });
27    }
28    sum
29}
30
31use std::{cmp::Ordering, collections::BinaryHeap};
32
33#[derive(Debug, Clone, Eq)]
34struct Vertex {
35    v: usize,
36    cost: i64,
37}
38impl Vertex {
39    #[inline]
40    pub fn new(v: usize, cost: i64) -> Self {
41        Vertex { v, cost }
42    }
43}
44impl PartialEq for Vertex {
45    fn eq(&self, other: &Vertex) -> bool {
46        self.cost.eq(&other.cost)
47    }
48}
49impl PartialOrd for Vertex {
50    fn partial_cmp(&self, other: &Vertex) -> Option<Ordering> {
51        Some(self.cmp(other))
52    }
53}
54impl Ord for Vertex {
55    fn cmp(&self, other: &Self) -> Ordering {
56        other.cost.cmp(&self.cost)
57    }
58}
59
60/// O(|E|log|V|)
61pub fn prim_heap(g: &[Vec<i64>]) -> Option<i64> {
62    let mut min_cost = vec![None; g.len()];
63    let mut heap = BinaryHeap::new();
64    heap.push(Vertex::new(0, 0));
65
66    let mut v_count = 0;
67
68    let mut total_cost = 0;
69    while let Some(Vertex { v, cost }) = heap.pop() {
70        if min_cost[v].is_some() {
71            continue;
72        }
73        total_cost += cost;
74        min_cost[v] = Some(total_cost);
75        v_count += 1;
76        if v_count == g.len() {
77            return Some(total_cost);
78        }
79
80        (0..g.len()).filter(|&i| g[v][i] != -1).for_each(|i| {
81            heap.push(Vertex::new(i, g[v][i]));
82        });
83    }
84    None
85}
86
87#[cfg(test)]
88mod tests {
89
90    use super::*;
91    #[test]
92    fn test_prim() {
93        let g = vec![
94            vec![-1, 2, 3, 1, -1],
95            vec![2, -1, -1, 4, -1],
96            vec![3, -1, -1, 1, 1],
97            vec![1, 4, 1, -1, 3],
98            vec![-1, -1, 1, 3, -1],
99        ];
100
101        assert_eq!(prim(&g), 5);
102    }
103    #[test]
104    fn test_prim_heap() {
105        let g = vec![
106            vec![-1, 2, 3, 1, -1],
107            vec![2, -1, -1, 4, -1],
108            vec![3, -1, -1, 1, 1],
109            vec![1, 4, 1, -1, 3],
110            vec![-1, -1, 1, 3, -1],
111        ];
112
113        assert_eq!(prim_heap(&g).unwrap(), 5);
114    }
115
116    #[test]
117    fn prim_heap_chooses_the_minimum_edge() {
118        let g = vec![vec![-1, 10, 1], vec![10, -1, 1], vec![1, 1, -1]];
119        assert_eq!(prim_heap(&g), Some(2));
120    }
121}
competitive_library/graph/minimum_spanning_tree_prim.rs

competitive_library/graph/
minimum_spanning_tree_prim.rs
