competitive/graph/
minimum_spanning_arborescence.rs

1use super::{EdgeListGraph, Group, MergingUnionFind, MonoidAct, PairingHeap, comparator::Less};
2
3impl EdgeListGraph {
4    /// tarjan
5    pub fn minimum_spanning_arborescence<G, F>(
6        &self,
7        root: usize,
8        weight: F,
9    ) -> Option<(G::T, Vec<usize>)>
10    where
11        G: Group<T: Ord>,
12        F: Fn(usize) -> G::T,
13    {
14        struct WeightAct<G>(std::marker::PhantomData<fn() -> G>);
15        impl<G> MonoidAct for WeightAct<G>
16        where
17            G: Group,
18        {
19            type Key = (G::T, usize);
20            type Act = G::T;
21            type ActMonoid = G;
22
23            fn act(x: &Self::Key, a: &Self::Act) -> Self::Key {
24                (G::operate(&x.0, a), x.1)
25            }
26
27            fn act_assign(x: &mut Self::Key, a: &Self::Act) {
28                x.0 = G::operate(&x.0, a);
29            }
30        }
31        let mut uf = MergingUnionFind::new_with_merger(
32            self.vertices_size(),
33            |_| PairingHeap::<(G::T, usize), Less, WeightAct<G>>::default(),
34            |x, y| x.append(y),
35        );
36        let mut state = vec![0; self.vertices_size()]; // 0: unprocessed, 1: in process, 2: completed
37        state[root] = 2;
38        for (id, &(_, to)) in self.edges().enumerate() {
39            uf.merge_data_mut(to).push((weight(id), id));
40        }
41        let mut paredge = vec![0; self.edges_size()];
42        let mut ord = vec![];
43        let mut leaf = vec![self.edges_size(); self.vertices_size()];
44        let mut cycle = 0usize;
45        let mut acc = G::unit();
46        for mut cur in self.vertices() {
47            if state[cur] != 0 {
48                continue;
49            }
50            let mut path = vec![];
51            let mut ch = vec![];
52            while state[cur] != 2 {
53                path.push(cur);
54                state[cur] = 1;
55                let (w, eid) = {
56                    match uf.merge_data_mut(cur).pop() {
57                        Some((w, eid)) => (w, eid),
58                        None => return None,
59                    }
60                };
61                uf.merge_data_mut(cur).apply_all(G::inverse(&w));
62                acc = G::operate(&acc, &w);
63                ord.push(eid);
64                let (u, v) = self[eid];
65                if leaf[v] >= self.edges_size() {
66                    leaf[v] = eid;
67                }
68                while cycle > 0 {
69                    paredge[ch.pop().unwrap()] = eid;
70                    cycle -= 1;
71                }
72                ch.push(eid);
73                if state[uf.find_root(u)] == 1 {
74                    while let Some(t) = path.pop() {
75                        state[t] = 2;
76                        cycle += 1;
77                        if !uf.unite(u, t) {
78                            break;
79                        }
80                    }
81                    state[uf.find_root(u)] = 1;
82                }
83                cur = uf.find_root(u);
84            }
85            for u in path.into_iter() {
86                state[u] = 2;
87            }
88        }
89        let mut tree = vec![root; self.vertices_size()];
90        let mut used = vec![false; self.edges_size()];
91        for eid in ord.into_iter().rev() {
92            if !used[eid] {
93                let (u, v) = self[eid];
94                tree[v] = u;
95                let mut x = leaf[v];
96                while x != eid {
97                    used[x] = true;
98                    x = paredge[x];
99                }
100            }
101        }
102        Some((acc, tree))
103    }
104}
competitive/graph/minimum_spanning_arborescence.rs

competitive/graph/
minimum_spanning_arborescence.rs
