use crate::{
graph::edge_weighted::{Edge, EdgeWeightedGraph},
priority_queue::{IndexPQ, PQ},
};
pub struct LazyPrimMST<'a> {
marked: Vec<bool>,
mst: Vec<&'a Edge>,
min_pq: PQ<&'a Edge>,
}
impl<'a> LazyPrimMST<'a> {
pub fn new(g: &'a EdgeWeightedGraph) -> Self {
let marked = vec![false; g.vertex_count()];
let mst = Vec::new();
let min_pq = PQ::new_min_pq();
let mut prim = LazyPrimMST {
marked,
mst,
min_pq,
};
prim.visit(g, 0);
while let Some(edge) = prim.min_pq.pop() {
let v = edge.either();
let w = edge.other(v);
if prim.marked[v] && prim.marked[w] {
continue;
}
prim.mst.push(edge);
for v in [v, w] {
if !prim.marked[v] {
prim.visit(g, v);
}
}
}
prim
}
fn visit(&mut self, g: &'a EdgeWeightedGraph, v: usize) {
self.marked[v] = true;
for e in g.adj(v) {
if !self.marked[e.other(v)] {
self.min_pq.push(e);
}
}
}
pub fn edges(&self) -> std::slice::Iter<&Edge> {
self.mst.iter()
}
pub fn weight(&self) -> f64 {
self.edges().map(|f| f.weight()).sum()
}
}
pub struct PrimMST<'a> {
edge_to: Vec<Option<&'a Edge>>,
dist_to: Vec<f64>,
marked: Vec<bool>,
min_pq: IndexPQ<f64>,
}
impl<'a> PrimMST<'a> {
pub fn new(g: &'a EdgeWeightedGraph) -> Self {
let mut edge_to = Vec::with_capacity(g.vertex_count());
for _ in 0..g.vertex_count() {
edge_to.push(None);
}
let dist_to = vec![f64::MAX; g.vertex_count()];
let marked = vec![false; g.vertex_count()];
let min_pq = IndexPQ::new_min_pq(g.vertex_count());
let mut prim = PrimMST {
edge_to,
dist_to,
marked,
min_pq,
};
prim.dist_to[0] = 0.0;
prim.min_pq.push(0, 0.0);
while let Some((k, _d)) = prim.min_pq.pop() {
prim.visit(g, k);
}
prim
}
fn visit(&mut self, g: &'a EdgeWeightedGraph, v: usize) {
self.marked[v] = true;
for e in g.adj(v) {
let w = e.other(v);
if self.marked[w] {
continue;
}
if e.weight() < self.dist_to[w] {
self.edge_to[w] = Some(e);
self.dist_to[w] = e.weight();
if self.min_pq.contains(w) {
self.min_pq.change(w, e.weight());
} else {
self.min_pq.push(w, e.weight());
}
}
}
}
pub fn edges(&self) -> Vec<&Edge> {
self.edge_to
.iter()
.skip(1)
.map(|f| f.as_deref().unwrap())
.collect()
}
pub fn weight(&self) -> f64 {
self.dist_to.iter().skip(1).sum()
}
}
#[cfg(test)]
mod test {
use crate::{
add_edge_weighted,
graph::edge_weighted::{
mst::prim::{LazyPrimMST, PrimMST},
EdgeWeightedGraph,
},
};
#[test]
fn test() {
let mut g = EdgeWeightedGraph::new(8);
add_edge_weighted!(g, 4, 5, 0.35);
add_edge_weighted!(g, 4, 7, 0.37);
add_edge_weighted!(g, 5, 7, 0.28);
add_edge_weighted!(g, 0, 7, 0.16);
add_edge_weighted!(g, 1, 5, 0.32);
add_edge_weighted!(g, 0, 4, 0.38);
add_edge_weighted!(g, 2, 3, 0.17);
add_edge_weighted!(g, 1, 7, 0.19);
add_edge_weighted!(g, 0, 2, 0.26);
add_edge_weighted!(g, 1, 2, 0.36);
add_edge_weighted!(g, 1, 3, 0.29);
add_edge_weighted!(g, 2, 7, 0.34);
add_edge_weighted!(g, 6, 2, 0.40);
add_edge_weighted!(g, 3, 6, 0.52);
add_edge_weighted!(g, 6, 0, 0.58);
add_edge_weighted!(g, 6, 4, 0.93);
let prim = LazyPrimMST::new(&g);
let mut gg = EdgeWeightedGraph::new(g.vertex_count());
for e in prim.edges() {
add_edge_weighted!(gg, e.v, e.w, e.weight);
}
let weight = prim.weight();
println!("weight:{}", weight as f32);
let prim = PrimMST::new(&g);
let mut gg = EdgeWeightedGraph::new(g.vertex_count());
for e in prim.edges() {
add_edge_weighted!(gg, e.v, e.w, e.weight);
}
println!("weight:{}", prim.weight() as f32);
assert_eq!(weight as f32, prim.weight() as f32);
}
}