1use params::*;
9use prelude::*;
10use props::Color;
11
12use fera_fun::first;
13
14use std::cmp::Ordering;
15use std::collections::BinaryHeap;
16use std::marker::PhantomData;
17use std::ops::DerefMut;
18
19pub trait Prim: Incidence {
20 fn prim<W, T>(
21 &self,
22 w: W,
23 ) -> PrimAlg<
24 &Self,
25 W,
26 NewVertexProp<Self, Color>,
27 NewVertexProp<Self, OptionEdge<Self>>,
28 Owned<PrimPriorityQueue<Self, T>>,
29 PhantomData<T>,
30 >
31 where
32 W: EdgePropGet<Self, T>,
33 T: Ord,
34 {
35 PrimAlg(
36 self,
37 w,
38 NewVertexProp(self, Color::White),
39 NewVertexProp(self, Self::edge_none()),
40 Owned(PrimPriorityQueue::<Self, T>::new()),
41 PhantomData,
42 )
43 }
44}
45
46impl<G: Incidence> Prim for G {}
47
48generic_struct! {
49 #[must_use]
50 pub struct PrimAlg(graph, weight, color, parent, queue, _marker)
51}
52
53impl<'a, G, W, C, P, Q, T> IntoIterator for PrimAlg<&'a G, W, C, P, Q, PhantomData<T>>
54where
55 G: Incidence + VertexList,
56 W: EdgePropGet<G, T>,
57 C: ParamDerefMut,
58 C::Target: VertexPropMut<G, Color>,
59 P: ParamDerefMut,
60 P::Target: VertexPropMut<G, OptionEdge<G>>,
61 Q: ParamDerefMut<Target = PrimPriorityQueue<G, T>>,
62 T: Ord + Default,
63{
64 type Item = Edge<G>;
65 type IntoIter = Iter<'a, G, W, C::Output, P::Output, Q::Output, T>;
66
67 fn into_iter(self) -> Self::IntoIter {
68 let PrimAlg(g, w, color, parent, queue, _) = self;
69 let mut color = color.build();
70 let mut queue = queue.build();
71 let v = first(g.vertices());
72 color[v] = Color::Gray;
73 queue.push(QueueItem::new(T::default(), v));
74 Iter {
75 g,
76 w,
77 color,
78 queue,
79 parent: parent.build(),
80 _marker: PhantomData,
81 }
82 }
83}
84
85pub struct Iter<'a, G, W, C, P, Q, T>
86where
87 G: 'a,
88{
89 g: &'a G,
90 w: W,
91 color: C,
92 parent: P,
93 queue: Q,
94 _marker: PhantomData<T>,
95}
96
97impl<'a, G, W, C, P, Q, T> Iterator for Iter<'a, G, W, C, P, Q, T>
98where
99 G: 'a + Incidence,
100 W: EdgePropGet<G, T>,
101 C: DerefMut,
102 C::Target: VertexPropMut<G, Color>,
103 P: DerefMut,
104 P::Target: VertexPropMut<G, OptionEdge<G>>,
105 Q: DerefMut<Target = PrimPriorityQueue<G, T>>,
106 T: Ord,
107{
108 type Item = Edge<G>;
109
110 fn next(&mut self) -> Option<Self::Item> {
111 while let Some(QueueItem { vertex: u, .. }) = self.queue.pop() {
112 if self.color[u] == Color::Black {
113 continue;
114 }
115 self.color[u] = Color::Black;
116 for e in self.g.out_edges(u) {
117 let v = self.g.target(e);
118 if self.color[v] == Color::Black {
119 continue;
120 }
121 if let Some(p) = self.parent[v].into_option() {
122 if self.w.get(p) < self.w.get(e) {
123 continue;
124 }
125 }
126 self.color[v] = Color::Gray;
127 self.parent[v] = e.into();
128 self.queue.push(QueueItem::new(self.w.get(e), v));
129 }
130 if let e @ Some(_) = self.parent[u].into_option() {
131 return e;
132 }
133 }
134 None
135 }
136}
137
138type PrimPriorityQueue<G, T> = BinaryHeap<QueueItem<T, Vertex<G>>>;
139
140pub struct QueueItem<A, B> {
141 prio: A,
142 vertex: B,
143}
144
145impl<A, B> QueueItem<A, B> {
146 pub fn new(prio: A, vertex: B) -> Self {
147 Self { prio, vertex }
148 }
149}
150
151impl<A: PartialEq, B> PartialEq for QueueItem<A, B> {
152 fn eq(&self, other: &Self) -> bool {
153 self.prio == other.prio
154 }
155}
156
157impl<A: Eq, B> Eq for QueueItem<A, B> {}
158
159impl<A: PartialOrd, B> PartialOrd for QueueItem<A, B> {
160 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
161 other.prio.partial_cmp(&self.prio)
162 }
163}
164
165impl<A: Ord, B> Ord for QueueItem<A, B> {
166 fn cmp(&self, other: &Self) -> Ordering {
167 other.prio.cmp(&self.prio)
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use super::Prim;
174 use fera_fun::vec;
175 use fun::sum_prop;
176 use prelude::*;
177
178 #[test]
179 fn mst() {
180 let g: StaticGraph = graph!(
181 5, (0, 4), (2, 3), (0, 1), (1, 4),
186 (1, 2), (2, 4),
188 (3, 4)
189 );
190 let mut weight = g.default_edge_prop(0usize);
191 for (e, &w) in g.edges().zip(&[1, 2, 3, 4, 5, 6, 7]) {
192 weight[e] = w;
193 }
194 let e = vec(g.edges());
195 let tree = vec(g.prim(&weight));
196 assert_eq!(11usize, sum_prop(&weight, &tree));
197 assert_eq!(vec![e[0], e[2], e[4], e[1]], tree);
198 }
199}