use super::reverse_path;
use crate::FxIndexMap;
use indexmap::map::Entry::{Occupied, Vacant};
use num_traits::Zero;
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::hash::Hash;
use std::usize;
pub fn dijkstra<N, C, FN, IN, FS>(
start: &N,
mut successors: FN,
mut success: FS,
) -> Option<(Vec<N>, C)>
where
N: Eq + Hash + Clone,
C: Zero + Ord + Copy,
FN: FnMut(&N) -> IN,
IN: IntoIterator<Item = (N, C)>,
FS: FnMut(&N) -> bool,
{
dijkstra_internal(start, &mut successors, &mut success)
}
pub(crate) fn dijkstra_internal<N, C, FN, IN, FS>(
start: &N,
successors: &mut FN,
success: &mut FS,
) -> Option<(Vec<N>, C)>
where
N: Eq + Hash + Clone,
C: Zero + Ord + Copy,
FN: FnMut(&N) -> IN,
IN: IntoIterator<Item = (N, C)>,
FS: FnMut(&N) -> bool,
{
let (parents, reached) = run_dijkstra(start, successors, success);
reached.map(|target| {
(
reverse_path(&parents, |&(p, _)| p, target),
parents.get_index(target).unwrap().1 .1,
)
})
}
pub fn dijkstra_all<N, C, FN, IN>(start: &N, successors: FN) -> HashMap<N, (N, C)>
where
N: Eq + Hash + Clone,
C: Zero + Ord + Copy,
FN: FnMut(&N) -> IN,
IN: IntoIterator<Item = (N, C)>,
{
dijkstra_partial(start, successors, |_| false).0
}
#[allow(clippy::missing_panics_doc)]
pub fn dijkstra_partial<N, C, FN, IN, FS>(
start: &N,
mut successors: FN,
mut stop: FS,
) -> (HashMap<N, (N, C)>, Option<N>)
where
N: Eq + Hash + Clone,
C: Zero + Ord + Copy,
FN: FnMut(&N) -> IN,
IN: IntoIterator<Item = (N, C)>,
FS: FnMut(&N) -> bool,
{
let (parents, reached) = run_dijkstra(start, &mut successors, &mut stop);
(
parents
.iter()
.skip(1)
.map(|(n, (p, c))| (n.clone(), (parents.get_index(*p).unwrap().0.clone(), *c))) .collect(),
reached.map(|i| parents.get_index(i).unwrap().0.clone()),
)
}
fn run_dijkstra<N, C, FN, IN, FS>(
start: &N,
successors: &mut FN,
stop: &mut FS,
) -> (FxIndexMap<N, (usize, C)>, Option<usize>)
where
N: Eq + Hash + Clone,
C: Zero + Ord + Copy,
FN: FnMut(&N) -> IN,
IN: IntoIterator<Item = (N, C)>,
FS: FnMut(&N) -> bool,
{
let mut to_see = BinaryHeap::new();
to_see.push(SmallestHolder {
cost: Zero::zero(),
index: 0,
});
let mut parents: FxIndexMap<N, (usize, C)> = FxIndexMap::default();
parents.insert(start.clone(), (usize::max_value(), Zero::zero()));
let mut target_reached = None;
while let Some(SmallestHolder { cost, index }) = to_see.pop() {
let successors = {
let (node, &(_, c)) = parents.get_index(index).unwrap();
if stop(node) {
target_reached = Some(index);
break;
}
if cost > c {
continue;
}
successors(node)
};
for (successor, move_cost) in successors {
let new_cost = cost + move_cost;
let n;
match parents.entry(successor) {
Vacant(e) => {
n = e.index();
e.insert((index, new_cost));
}
Occupied(mut e) => {
if e.get().1 > new_cost {
n = e.index();
e.insert((index, new_cost));
} else {
continue;
}
}
}
to_see.push(SmallestHolder {
cost: new_cost,
index: n,
});
}
}
(parents, target_reached)
}
#[allow(clippy::implicit_hasher)]
pub fn build_path<N, C>(target: &N, parents: &HashMap<N, (N, C)>) -> Vec<N>
where
N: Eq + Hash + Clone,
{
let mut rev = vec![target.clone()];
let mut next = target.clone();
while let Some((parent, _)) = parents.get(&next) {
rev.push(parent.clone());
next = parent.clone();
}
rev.reverse();
rev
}
struct SmallestHolder<K> {
cost: K,
index: usize,
}
impl<K: PartialEq> PartialEq for SmallestHolder<K> {
fn eq(&self, other: &Self) -> bool {
self.cost == other.cost
}
}
impl<K: PartialEq> Eq for SmallestHolder<K> {}
impl<K: Ord> PartialOrd for SmallestHolder<K> {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl<K: Ord> Ord for SmallestHolder<K> {
fn cmp(&self, other: &Self) -> Ordering {
other.cost.cmp(&self.cost)
}
}