haar_lib/graph/
chu_liu_edmonds.rs

1//! 有向グラフ上の最小有向全域木を求める
2//! # Problems
3//! - [AOJ GRL_2_B](https://onlinejudge.u-aizu.ac.jp/problems/GRL_2_B)
4
5use crate::graph::*;
6use std::collections::VecDeque;
7use std::ops::Sub;
8
9type Edge_<'a, T, E> = (usize, usize, T, &'a E);
10
11/// 有向グラフ上の最小有向全域木を求める
12///
13/// **Time complexity** $O(VE)$
14pub fn chu_liu_edmonds<E: EdgeTrait>(g: &Graph<Directed, E>, root: usize) -> Vec<&E>
15where
16    E::Weight: Ord + Copy + Sub<Output = E::Weight>,
17{
18    let n = g.len();
19    let mut rg = vec![vec![]; n];
20
21    for e in g.nodes_iter().flat_map(|v| &v.edges) {
22        rg[e.to()].push((e.to(), e.from(), e.weight(), e));
23    }
24
25    let res = rec(rg, root);
26    res.into_iter().filter_map(|e| e.map(|e| e.3)).collect()
27}
28
29fn rec<T, E>(mut g: Vec<Vec<Edge_<T, E>>>, root: usize) -> Vec<Option<Edge_<T, E>>>
30where
31    T: Ord + Copy + Sub<Output = T>,
32    E: EdgeTrait<Weight = T>,
33{
34    let n = g.len();
35
36    let mut in_edges = vec![None; n];
37    let mut out_count = vec![0; n];
38
39    for (i, es) in g.iter().enumerate() {
40        if i != root {
41            let &res @ (from, to, ..) = es.iter().min_by(|x, y| x.2.cmp(&y.2)).unwrap();
42
43            in_edges[from] = Some(res);
44            out_count[to] += 1;
45        }
46    }
47
48    let mut q: VecDeque<_> = out_count
49        .iter()
50        .enumerate()
51        .filter_map(|(i, &x)| (x == 0).then_some(i))
52        .collect();
53
54    while let Some(i) = q.pop_front() {
55        if let Some((_, to, ..)) = in_edges[i] {
56            out_count[to] -= 1;
57
58            if out_count[to] == 0 {
59                q.push_back(to);
60            }
61        }
62    }
63
64    let cycles = {
65        let temp = out_count
66            .into_iter()
67            .enumerate()
68            .filter_map(|(i, c)| (c != 0).then_some(i));
69
70        let mut ret = vec![];
71        let mut check = vec![false; n];
72
73        for i in temp {
74            if check[i] {
75                continue;
76            }
77
78            let mut cur = i;
79            let mut cycle = vec![];
80
81            while !check[cur] {
82                check[cur] = true;
83                cycle.push(cur);
84                cur = in_edges[cur].unwrap().1;
85            }
86
87            ret.push(cycle);
88        }
89
90        ret
91    };
92
93    if !cycles.is_empty() {
94        let mut s = vec![0; n];
95        let mut in_cycle = vec![false; n];
96        let mut groups = vec![];
97
98        for cycle in cycles {
99            for &i in &cycle {
100                let &c = g[i].iter().min_by(|x, y| x.2.cmp(&y.2)).unwrap();
101
102                for e in g[i].iter_mut() {
103                    e.2 = e.2 - c.2;
104                }
105
106                in_cycle[i] = true;
107            }
108
109            groups.push(cycle);
110        }
111
112        for (i, x) in in_cycle.into_iter().enumerate() {
113            if !x {
114                groups.push(vec![i]);
115            }
116        }
117
118        let size = groups.len();
119
120        for (i, xs) in groups.iter().enumerate() {
121            for &x in xs {
122                s[x] = i;
123            }
124        }
125
126        let mut sg = vec![vec![]; size];
127        let root = s[root];
128
129        for &(from, to, weight, e) in g.iter().flatten() {
130            if s[from] != s[to] {
131                sg[s[from]].push((s[from], s[to], weight, e));
132            }
133        }
134
135        let res = rec(sg, root);
136
137        for (i, e) in res.into_iter().enumerate() {
138            let p = if let Some(e) = e { e.1 } else { continue };
139
140            let c = groups[i]
141                .iter()
142                .flat_map(|&x| g[x].iter())
143                .filter(|e| s[e.1] == p);
144
145            let &res @ (from, ..) = c.min_by(|x, y| x.2.cmp(&y.2)).unwrap();
146
147            in_edges[from] = Some(res);
148        }
149    }
150
151    in_edges
152}
haar_lib/graph/chu_liu_edmonds.rs

haar_lib/graph/
chu_liu_edmonds.rs
