petgraph/algo/
johnson.rs

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

petgraph/algo/
johnson.rs
