competitive/graph/
minimum_spanning_tree.rs

1use super::EdgeListGraph;
2use crate::algebra::Group;
3use crate::data_structure::{MergingUnionFind, UnionFind};
4
5#[codesnip::entry("minimum_spanning_tree", include("EdgeListGraph", "UnionFind"))]
6impl EdgeListGraph {
7    pub fn minimum_spanning_tree<T>(&self, weight: impl Fn(&usize) -> T) -> Vec<bool>
8    where
9        T: Ord,
10    {
11        let mut idx: Vec<_> = (0..self.edges_size()).collect();
12        idx.sort_by_key(weight);
13        let mut uf = UnionFind::new(self.vertices_size());
14        let mut res = vec![false; self.edges_size()];
15        for eid in idx.into_iter() {
16            let (u, v) = self[eid];
17            res[eid] = uf.unite(u, v);
18        }
19        res
20    }
21}
22
23#[codesnip::entry(
24    "minimum_spanning_arborescence",
25    include("algebra", "EdgeListGraph", "UnionFind")
26)]
27impl EdgeListGraph {
28    /// tarjan
29    pub fn minimum_spanning_arborescence<G, F>(
30        &self,
31        root: usize,
32        weight: F,
33    ) -> Option<(G::T, Vec<usize>)>
34    where
35        G: Group,
36        G::T: Ord,
37        F: Fn(usize) -> G::T,
38    {
39        use std::{cmp::Reverse, collections::BinaryHeap};
40        let mut uf = MergingUnionFind::new_with_merger(
41            self.vertices_size(),
42            |_| (BinaryHeap::new(), G::unit()),
43            |x, y| {
44                let ny = G::rinv_operate(&y.1, &x.1);
45                x.0.extend(
46                    (y.0)
47                        .drain()
48                        .map(|(Reverse(ref w), i)| (Reverse(G::operate(w, &ny)), i)),
49                )
50            },
51        );
52        let mut state = vec![0; self.vertices_size()]; // 0: unprocessed, 1: in process, 2: completed
53        state[root] = 2;
54        for (id, &(_, to)) in self.edges().enumerate() {
55            uf.merge_data_mut(to).0.push((Reverse(weight(id)), id));
56        }
57        let mut paredge = vec![0; self.edges_size()];
58        let mut ord = vec![];
59        let mut leaf = vec![self.edges_size(); self.vertices_size()];
60        let mut cycle = 0usize;
61        let mut acc = G::unit();
62        for mut cur in self.vertices() {
63            if state[cur] != 0 {
64                continue;
65            }
66            let mut path = vec![];
67            let mut ch = vec![];
68            while state[cur] != 2 {
69                path.push(cur);
70                state[cur] = 1;
71                let (w, eid) = {
72                    let (heap, lazy) = &mut uf.merge_data_mut(cur);
73                    match heap.pop() {
74                        Some((Reverse(w), eid)) => (G::operate(&w, lazy), eid),
75                        None => return None,
76                    }
77                };
78                {
79                    let curw = &mut uf.merge_data_mut(cur).1;
80                    *curw = G::rinv_operate(curw, &w);
81                }
82                acc = G::operate(&acc, &w);
83                ord.push(eid);
84                let (u, v) = self[eid];
85                if leaf[v] >= self.edges_size() {
86                    leaf[v] = eid;
87                }
88                while cycle > 0 {
89                    paredge[ch.pop().unwrap()] = eid;
90                    cycle -= 1;
91                }
92                ch.push(eid);
93                if state[uf.find_root(u)] == 1 {
94                    while let Some(t) = path.pop() {
95                        state[t] = 2;
96                        cycle += 1;
97                        if !uf.unite(u, t) {
98                            break;
99                        }
100                    }
101                    state[uf.find_root(u)] = 1;
102                }
103                cur = uf.find_root(u);
104            }
105            for u in path.into_iter() {
106                state[u] = 2;
107            }
108        }
109        let mut tree = vec![root; self.vertices_size()];
110        let mut used = vec![false; self.edges_size()];
111        for eid in ord.into_iter().rev() {
112            if !used[eid] {
113                let (u, v) = self[eid];
114                tree[v] = u;
115                let mut x = leaf[v];
116                while x != eid {
117                    used[x] = true;
118                    x = paredge[x];
119                }
120            }
121        }
122        Some((acc, tree))
123    }
124}
competitive/graph/minimum_spanning_tree.rs

competitive/graph/
minimum_spanning_tree.rs
