competitive/graph/minimum_spanning_tree.rs
1use super::{EdgeListGraph, UnionFind};
2
3impl EdgeListGraph {
4 pub fn minimum_spanning_tree<T>(&self, weight: impl Fn(&usize) -> T) -> Vec<bool>
5 where
6 T: Ord,
7 {
8 let mut edges: Vec<_> = (0..self.edges_size()).collect();
9 edges.sort_unstable_by_key(weight);
10 self.minimum_spanning_tree_from_sorted_edges(edges)
11 }
12
13 /// Runs Kruskal's algorithm on edge IDs sorted by nondecreasing weight.
14 pub fn minimum_spanning_tree_from_sorted_edges(
15 &self,
16 edges: impl IntoIterator<Item = usize>,
17 ) -> Vec<bool> {
18 let mut uf = UnionFind::new(self.vertices_size());
19 let mut res = vec![false; self.edges_size()];
20 let mut selected = 0;
21 for eid in edges {
22 let (u, v) = self[eid];
23 res[eid] = uf.unite(u, v);
24 if res[eid] {
25 selected += 1;
26 if selected + 1 == self.vertices_size() {
27 break;
28 }
29 }
30 }
31 res
32 }
33}