petgraph/algo/
johnson.rs

1//! Johnson's algorithm implementation.
2use alloc::{vec, vec::Vec};
3use core::hash::Hash;
4use core::ops::Sub;
5
6use hashbrown::HashMap;
7
8use super::{dijkstra, spfa::spfa_loop};
9pub use super::{BoundedMeasure, NegativeCycle};
10use crate::visit::{EdgeRef, IntoEdges, IntoNodeIdentifiers, NodeIndexable, Visitable};
11
12#[cfg(feature = "rayon")]
13use core::marker::{Send, Sync};
14
15/// \[Generic\] [Johnson algorithm][johnson] for all pairs shortest path problem.
16///
17/// Сompute the lengths of shortest paths in a weighted graph with
18/// positive or negative edge weights, but no negative cycles.
19///
20/// The time complexity of this implementation is **O(|V||E|log(|V|) + |V|²*log(|V|))**,
21/// which is faster than [`floyd_warshall`](fn@crate::algo::floyd_warshall) on sparse graphs and slower on dense ones.
22///
23/// If you are working with a sparse graph that is guaranteed to have no negative weights,
24/// it's preferable to run [`dijkstra`](fn@crate::algo::dijkstra) several times.
25///
26/// There is also a parallel implementation `parallel_johnson`, available under the `rayon` feature.
27///
28/// ## Arguments
29/// * `graph`: weighted graph.
30/// * `edge_cost`: closure that returns cost of a particular edge.
31///
32/// ## Returns
33/// * `Err`: if graph contains negative cycle.
34/// * `Ok`: `HashMap` of shortest distances.
35///
36/// # Complexity
37/// * Time complexity: **O(|V||E|log(|V|) + |V|²log(|V|))** since the implementation is based on [`dijkstra`](fn@crate::algo::dijkstra).
38/// * Auxiliary space: **O(|V|² + |V||E|)**.
39///
40/// where **|V|** is the number of nodes and **|E|** is the number of edges.
41///
42/// [johnson]: https://en.wikipedia.org/wiki/Johnson%27s_algorithm
43///
44/// # Examples
45///
46/// ```
47/// use petgraph::{prelude::*, Graph, Directed};
48/// use petgraph::algo::johnson;
49/// use std::collections::HashMap;
50///
51/// let mut graph: Graph<(), i32, Directed> = Graph::new();
52/// let a = graph.add_node(());
53/// let b = graph.add_node(());
54/// let c = graph.add_node(());
55/// let d = graph.add_node(());
56///
57/// graph.extend_with_edges(&[
58///    (a, b, 1),
59///    (a, c, 4),
60///    (a, d, 10),
61///    (b, c, 2),
62///    (b, d, 2),
63///    (c, d, 2)
64/// ]);
65///
66/// //     ----- b --------
67/// //    |      ^         | 2
68/// //    |    1 |    4    v
69/// //  2 |      a ------> c
70/// //    |   10 |         | 2
71/// //    |      v         v
72/// //     --->  d <-------
73///
74/// let expected_res: HashMap<(NodeIndex, NodeIndex), i32> = [
75///    ((a, a), 0), ((a, b), 1), ((a, c), 3), ((a, d), 3),
76///    ((b, b), 0), ((b, c), 2), ((b, d), 2),
77///    ((c, c), 0), ((c, d), 2),
78///    ((d, d), 0),
79/// ].iter().cloned().collect();
80///
81///
82/// let res = johnson(&graph, |edge| {
83///     *edge.weight()
84/// }).unwrap();
85///
86/// let nodes = [a, b, c, d];
87/// for node1 in &nodes {
88///     for node2 in &nodes {
89///         assert_eq!(res.get(&(*node1, *node2)), expected_res.get(&(*node1, *node2)));
90///     }
91/// }
92/// ```
93#[allow(clippy::type_complexity)]
94pub fn johnson<G, F, K>(
95    graph: G,
96    mut edge_cost: F,
97) -> Result<HashMap<(G::NodeId, G::NodeId), K>, NegativeCycle>
98where
99    G: IntoEdges + IntoNodeIdentifiers + NodeIndexable + Visitable,
100    G::NodeId: Eq + Hash,
101    F: FnMut(G::EdgeRef) -> K,
102    K: BoundedMeasure + Copy + Sub<K, Output = K>,
103{
104    let reweight = johnson_reweight(graph, &mut edge_cost)?;
105    let reweight = reweight.as_slice();
106
107    let node_bound = graph.node_bound();
108    let ix = |i| graph.to_index(i);
109
110    let mut distance_map: HashMap<(G::NodeId, G::NodeId), K> =
111        HashMap::with_capacity(node_bound * node_bound);
112
113    // Reweight edges.
114    let mut new_cost = |edge: G::EdgeRef| {
115        let (sum, _overflow) = edge_cost(edge).overflowing_add(reweight[ix(edge.source())]);
116        debug_assert!(!_overflow);
117        sum - reweight[ix(edge.target())]
118    };
119
120    // Run Dijkstra's algorithm from each node.
121    for source in graph.node_identifiers() {
122        for (target, dist) in dijkstra(graph, source, None, &mut new_cost) {
123            distance_map.insert(
124                (source, target),
125                dist + reweight[ix(target)] - reweight[ix(source)],
126            );
127        }
128    }
129
130    Ok(distance_map)
131}
132
133/// \[Generic\] [Johnson algorithm][johnson]
134/// implementation for all pairs shortest path problem,
135/// parallelizing the [`dijkstra`](fn@crate::algo::dijkstra) calls with `rayon`.
136///
137/// Сompute the lengths of shortest paths in a weighted graph with
138/// positive or negative edge weights, but no negative cycles.
139///
140/// If you are working with a sparse graph that is guaranteed to have no negative weights,
141/// it's preferable to run [`dijkstra`](fn@crate::algo::dijkstra) several times in parallel.
142///
143/// ## Arguments
144/// * `graph`: weighted graph.
145/// * `edge_cost`: closure that returns cost of a particular edge.
146///
147/// ## Returns
148/// * `Err`: if graph contains negative cycle.
149/// * `Ok`: `HashMap` of shortest distances.
150///
151/// [johnson]: https://en.wikipedia.org/wiki/Johnson%27s_algorithm
152///
153/// # Examples
154///
155/// ```
156/// use petgraph::{prelude::*, Graph, Directed};
157/// use petgraph::algo::parallel_johnson;
158/// use std::collections::HashMap;
159///
160/// let mut graph: Graph<(), i32, Directed> = Graph::new();
161/// let a = graph.add_node(());
162/// let b = graph.add_node(());
163/// let c = graph.add_node(());
164/// let d = graph.add_node(());
165///
166/// graph.extend_with_edges(&[
167///    (a, b, 1),
168///    (a, c, 4),
169///    (a, d, 10),
170///    (b, c, 2),
171///    (b, d, 2),
172///    (c, d, 2)
173/// ]);
174///
175/// //     ----- b --------
176/// //    |      ^         | 2
177/// //    |    1 |    4    v
178/// //  2 |      a ------> c
179/// //    |   10 |         | 2
180/// //    |      v         v
181/// //     --->  d <-------
182///
183/// let expected_res: HashMap<(NodeIndex, NodeIndex), i32> = [
184///    ((a, a), 0), ((a, b), 1), ((a, c), 3), ((a, d), 3),
185///    ((b, b), 0), ((b, c), 2), ((b, d), 2),
186///    ((c, c), 0), ((c, d), 2),
187///    ((d, d), 0),
188/// ].iter().cloned().collect();
189///
190///
191/// let res = parallel_johnson(&graph, |edge| {
192///     *edge.weight()
193/// }).unwrap();
194///
195/// let nodes = [a, b, c, d];
196/// for node1 in &nodes {
197///     for node2 in &nodes {
198///         assert_eq!(res.get(&(*node1, *node2)), expected_res.get(&(*node1, *node2)));
199///     }
200/// }
201/// ```
202#[cfg(feature = "rayon")]
203#[allow(clippy::type_complexity)]
204pub fn parallel_johnson<G, F, K>(
205    graph: G,
206    mut edge_cost: F,
207) -> Result<HashMap<(G::NodeId, G::NodeId), K>, NegativeCycle>
208where
209    G: IntoEdges + IntoNodeIdentifiers + NodeIndexable + Visitable + Sync,
210    G::NodeId: Eq + Hash + Send,
211    F: Fn(G::EdgeRef) -> K + Sync,
212    K: BoundedMeasure + Copy + Sub<K, Output = K> + Send + Sync,
213{
214    use rayon::iter::{IntoParallelIterator, ParallelIterator};
215
216    let reweight = johnson_reweight(graph, &mut edge_cost)?;
217    let reweight = reweight.as_slice();
218
219    let node_bound = graph.node_bound();
220    let ix = |i| graph.to_index(i);
221
222    // Reweight edges.
223    let new_cost = |edge: G::EdgeRef| {
224        let (sum, _overflow) = edge_cost(edge).overflowing_add(reweight[ix(edge.source())]);
225        debug_assert!(!_overflow);
226        sum - reweight[ix(edge.target())]
227    };
228
229    // Run Dijkstra's algorithm from each node.
230    let distance_map = (0..node_bound)
231        .into_par_iter()
232        .flat_map_iter(|s| {
233            let source = graph.from_index(s);
234
235            dijkstra(graph, source, None, new_cost)
236                .into_iter()
237                .map(move |(target, dist)| {
238                    (
239                        (source, target),
240                        dist + reweight[ix(target)] - reweight[ix(source)],
241                    )
242                })
243        })
244        .collect::<HashMap<(G::NodeId, G::NodeId), K>>();
245
246    Ok(distance_map)
247}
248
249/// Add a virtual node to the graph with oriented edges with zero weight
250/// to all other vertices, and then run SPFA from it.
251/// The found distances will be used to change the edge weights in Dijkstra's
252/// algorithm to make them non-negative.
253fn johnson_reweight<G, F, K>(graph: G, mut edge_cost: F) -> Result<Vec<K>, NegativeCycle>
254where
255    G: IntoEdges + IntoNodeIdentifiers + NodeIndexable + Visitable,
256    G::NodeId: Eq + Hash,
257    F: FnMut(G::EdgeRef) -> K,
258    K: BoundedMeasure + Copy + Sub<K, Output = K>,
259{
260    let node_bound = graph.node_bound();
261
262    let reweight = vec![K::default(); node_bound];
263
264    // Queue of vertices capable of relaxation of the found shortest distances.
265    let mut queue: Vec<G::NodeId> = Vec::with_capacity(node_bound);
266
267    // Adding all vertices to the queue is the same as starting the algorithm from a virtual node.
268    queue.extend(graph.node_identifiers());
269    let in_queue = vec![true; node_bound];
270
271    spfa_loop(graph, reweight, None, queue, in_queue, &mut edge_cost).map(|(dists, _)| dists)
272}
petgraph/algo/johnson.rs

petgraph/algo/
johnson.rs
