1use params::*;
10use prelude::*;
11use unionfind::{NewUnionFind, UnionFind, WithUnionFind};
12
13use fera_fun::vec;
14
15use std::ops::DerefMut;
16
17pub trait Visitor<G>
18where
19 G: WithEdge + WithUnionFind,
20{
21 fn accept(&mut self, g: &G, e: Edge<G>, ds: &mut UnionFind<G>) -> bool;
22
23 #[allow(unused_variables)]
24 fn after_union(&mut self, g: &G, e: Edge<G>, ds: &mut UnionFind<G>) {}
25}
26
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28pub struct AcceptAll;
29
30impl<G> Visitor<G> for AcceptAll
31where
32 G: WithEdge + WithUnionFind,
33{
34 fn accept(&mut self, _g: &G, _e: Edge<G>, _ds: &mut UnionFind<G>) -> bool {
35 true
36 }
37}
38
39impl<F, G> Visitor<G> for F
40where
41 G: WithEdge + WithUnionFind,
42 F: FnMut(&G, Edge<G>, &mut UnionFind<G>) -> bool,
43{
44 fn accept(&mut self, g: &G, e: Edge<G>, ds: &mut UnionFind<G>) -> bool {
45 self(g, e, ds)
46 }
47}
48
49pub struct Iter<'a, G: 'a, E, V, U> {
50 g: &'a G,
51 edges: E,
52 visitor: V,
53 ds: U,
54}
55
56impl<'a, G, E, V, U> Iterator for Iter<'a, G, E, V, U>
57where
58 G: 'a + WithUnionFind,
59 E: Iterator,
60 E::Item: IntoOwned<Edge<G>>,
61 V: Visitor<G>,
62 U: DerefMut<Target = UnionFind<G>>,
63{
64 type Item = Edge<G>;
65
66 fn next(&mut self) -> Option<Edge<G>> {
67 if self.ds.num_sets() > 1 {
68 for e in self.edges.by_ref().map(|e| e.into_owned()) {
69 let (u, v) = self.g.ends(e);
70 if !self.ds.in_same_set(u, v) && self.visitor.accept(self.g, e, &mut self.ds) {
71 self.ds.union(u, v);
72 self.visitor.after_union(self.g, e, &mut self.ds);
73 return Some(e);
74 }
75 }
76 }
77 None
78 }
79}
80
81pub trait Kruskal: WithUnionFind {
82 fn kruskal_mst<T, W>(
83 &self,
84 weight: W,
85 ) -> KruskalAlg<&Self, Vec<Edge<Self>>, AcceptAll, NewUnionFind<Self>>
86 where
87 W: EdgePropGet<Self, T>,
88 T: Ord,
89 {
90 self.kruskal().weight(weight)
91 }
92
93 fn kruskal(&self) -> KruskalAlg<&Self, AllEdges<Self>, AcceptAll, NewUnionFind<Self>> {
94 KruskalAlg(self, AllEdges(self), AcceptAll, NewUnionFind(self))
95 }
96}
97
98impl<G: WithUnionFind> Kruskal for G {}
99
100generic_struct! {
101 #[must_use]
102 pub struct KruskalAlg(graph, edges, visitor, unionfind)
103}
104
105impl<'a, G, E, V, U> KruskalAlg<&'a G, E, V, U>
106where
107 G: WithUnionFind,
108{
109 pub fn weight<W, T>(self, w: W) -> KruskalAlg<&'a G, Vec<Edge<G>>, V, U>
110 where
111 W: EdgePropGet<G, T>,
112 T: Ord,
113 {
114 let edges = vec(self.0.edges()).sorted_by_prop(&w);
115 self.edges(edges)
116 }
117}
118
119impl<'a, G, E, V, U> IntoIterator for KruskalAlg<&'a G, E, V, U>
120where
121 G: WithUnionFind,
122 E: IntoIterator,
123 E::Item: IntoOwned<Edge<G>>,
124 V: Visitor<G>,
125 U: ParamDerefMut<Target = UnionFind<G>>,
126{
127 type Item = Edge<G>;
128 type IntoIter = Iter<'a, G, E::IntoIter, V, U::Output>;
129
130 fn into_iter(self) -> Self::IntoIter {
131 let KruskalAlg(g, edges, visitor, ds) = self;
132 Iter {
133 g,
134 visitor,
135 edges: edges.into_iter(),
136 ds: ds.build(),
137 }
138 }
139}
140
141#[cfg(test)]
142mod tests {
143 use super::Kruskal;
144 use fera_fun::vec;
145 use prelude::*;
146
147 #[test]
148 fn kruskal_mst() {
149 let g: StaticGraph = graph!(
150 5,
151 (0, 4), (2, 3), (0, 1), (1, 4),
155 (1, 2), (2, 4),
157 (3, 4),
158 );
159 let mut weight = g.default_edge_prop(0usize);
160 for (e, &w) in g.edges().zip(&[1, 2, 3, 4, 5, 6, 7]) {
161 weight[e] = w;
162 }
163 let e = vec(g.edges());
164 assert_eq!(vec![e[0], e[1], e[2], e[4]], vec(g.kruskal_mst(&weight)));
165 }
166}
167
168