use super::reverse_path;
use crate::FxIndexMap;
use indexmap::map::Entry::{Occupied, Vacant};
use num_traits::Zero;
use rustc_hash::{FxHashMap, FxHashSet};
use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap};
use std::hash::Hash;
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
}
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, 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)
}
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)
}
}
pub struct DijkstraReachable<N, C, FN> {
to_see: BinaryHeap<SmallestHolder<C>>,
seen: FxHashSet<usize>,
parents: FxIndexMap<N, (usize, C)>,
total_costs: FxHashMap<N, C>,
successors: FN,
}
#[derive(Debug, Hash, PartialEq, Eq, Clone)]
pub struct DijkstraReachableItem<N, C> {
pub node: N,
pub parent: Option<N>,
pub total_cost: C,
}
impl<N, C, FN, IN> Iterator for DijkstraReachable<N, C, FN>
where
N: Eq + Hash + Clone,
C: Zero + Ord + Copy + Hash,
FN: FnMut(&N) -> IN,
IN: IntoIterator<Item = (N, C)>,
{
type Item = DijkstraReachableItem<N, C>;
fn next(&mut self) -> Option<Self::Item> {
while let Some(SmallestHolder { cost, index }) = self.to_see.pop() {
if !self.seen.insert(index) {
continue;
}
let item;
let successors = {
let (node, (parent_index, _)) = self.parents.get_index(index).unwrap();
let total_cost = self.total_costs[node];
item = Some(DijkstraReachableItem {
node: node.clone(),
parent: self.parents.get_index(*parent_index).map(|x| x.0.clone()),
total_cost,
});
(self.successors)(node)
};
for (successor, move_cost) in successors {
let new_cost = cost + move_cost;
let n;
match self.parents.entry(successor.clone()) {
Vacant(e) => {
n = e.index();
e.insert((index, new_cost));
self.total_costs.insert(successor.clone(), new_cost);
}
Occupied(mut e) => {
if e.get().1 > new_cost {
n = e.index();
e.insert((index, new_cost));
self.total_costs.insert(successor.clone(), new_cost);
} else {
continue;
}
}
}
self.to_see.push(SmallestHolder {
cost: new_cost,
index: n,
});
}
return item;
}
None
}
}
pub fn dijkstra_reach<N, C, FN, IN>(start: &N, successors: FN) -> DijkstraReachable<N, C, FN>
where
N: Eq + Hash + Clone,
C: Zero + Ord + Copy,
FN: FnMut(&N) -> IN,
IN: IntoIterator<Item = (N, C)>,
{
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, Zero::zero()));
let mut total_costs = FxHashMap::default();
total_costs.insert(start.clone(), Zero::zero());
let seen = FxHashSet::default();
DijkstraReachable {
to_see,
seen,
parents,
total_costs,
successors,
}
}