dijkstra's algorithm
Ziyang Hu
2 years ago
parent
00fff0ce5a
commit
0e3c9507e4
@ -0,0 +1,187 @@
|
||||
use std::cmp::Ordering;
|
||||
use std::collections::{BTreeMap, BTreeSet, BinaryHeap};
|
||||
|
||||
use anyhow::{anyhow, bail, Result};
|
||||
use either::{Left, Right};
|
||||
use itertools::Itertools;
|
||||
|
||||
use crate::algo::AlgoImpl;
|
||||
use crate::data::expr::Expr;
|
||||
use crate::data::program::{MagicAlgoRuleArg, MagicSymbol};
|
||||
use crate::data::symb::Symbol;
|
||||
use crate::data::tuple::Tuple;
|
||||
use crate::data::value::DataValue;
|
||||
use crate::runtime::derived::DerivedRelStore;
|
||||
use crate::runtime::transact::SessionTx;
|
||||
|
||||
pub(crate) struct ShortestPathDijkstra;
|
||||
|
||||
impl AlgoImpl for ShortestPathDijkstra {
|
||||
fn run(
|
||||
&mut self,
|
||||
tx: &mut SessionTx,
|
||||
rels: &[MagicAlgoRuleArg],
|
||||
opts: &BTreeMap<Symbol, Expr>,
|
||||
stores: &BTreeMap<MagicSymbol, DerivedRelStore>,
|
||||
out: &DerivedRelStore,
|
||||
) -> Result<()> {
|
||||
let edges = rels
|
||||
.get(0)
|
||||
.ok_or_else(|| anyhow!("'shortest_path_dijkstra' requires edges relation"))?;
|
||||
let starting = rels.get(1).ok_or_else(|| {
|
||||
anyhow!("'shortest_path_dijkstra' requires starting relation as second argument")
|
||||
})?;
|
||||
let termination = rels.get(2);
|
||||
let allow_negative_edges = match opts.get(&Symbol::from("allow_negative_edges")) {
|
||||
None => false,
|
||||
Some(Expr::Const(DataValue::Bool(b))) => *b,
|
||||
Some(v) => bail!("option 'allow_negative_edges' for 'shortest_path_dijkstra' requires a boolean, got {:?}", v)
|
||||
};
|
||||
let undirected = match opts.get(&Symbol::from("undirected")) {
|
||||
None => false,
|
||||
Some(Expr::Const(DataValue::Bool(b))) => *b,
|
||||
Some(v) => bail!(
|
||||
"option 'undirected' for 'shortest_path_dijkstra' requires a boolean, got {:?}",
|
||||
v
|
||||
),
|
||||
};
|
||||
|
||||
let (graph, indices, inv_indices) =
|
||||
edges.convert_edge_to_weighted_graph(undirected, allow_negative_edges, tx, stores)?;
|
||||
|
||||
let mut starting_nodes = BTreeSet::new();
|
||||
for tuple in starting.iter(tx, stores)? {
|
||||
let tuple = tuple?;
|
||||
let node = tuple
|
||||
.0
|
||||
.get(0)
|
||||
.ok_or_else(|| anyhow!("node relation too short"))?;
|
||||
if let Some(idx) = inv_indices.get(node) {
|
||||
starting_nodes.insert(*idx);
|
||||
}
|
||||
}
|
||||
let termination_nodes = match termination {
|
||||
None => None,
|
||||
Some(t) => {
|
||||
let mut tn = BTreeSet::new();
|
||||
for tuple in t.iter(tx, stores)? {
|
||||
let tuple = tuple?;
|
||||
let node = tuple
|
||||
.0
|
||||
.get(0)
|
||||
.ok_or_else(|| anyhow!("node relation too short"))?;
|
||||
if let Some(idx) = inv_indices.get(node) {
|
||||
tn.insert(*idx);
|
||||
}
|
||||
}
|
||||
Some(tn)
|
||||
}
|
||||
};
|
||||
|
||||
for start in starting_nodes {
|
||||
let res = dijkstra(&graph, start, &termination_nodes);
|
||||
for (target, cost, path) in res {
|
||||
let t = vec![
|
||||
indices[start].clone(),
|
||||
indices[target].clone(),
|
||||
DataValue::from(cost),
|
||||
DataValue::List(path.into_iter().map(|u| indices[u].clone()).collect_vec()),
|
||||
];
|
||||
out.put(Tuple(t), 0)
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(PartialEq)]
|
||||
struct HeapState {
|
||||
cost: f64,
|
||||
node: usize,
|
||||
}
|
||||
|
||||
impl PartialOrd for HeapState {
|
||||
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
|
||||
Some(self.cmp(other))
|
||||
}
|
||||
}
|
||||
|
||||
impl Ord for HeapState {
|
||||
fn cmp(&self, other: &Self) -> Ordering {
|
||||
self.cost
|
||||
.total_cmp(&other.cost)
|
||||
.reverse()
|
||||
.then_with(|| self.node.cmp(&other.node))
|
||||
}
|
||||
}
|
||||
|
||||
impl Eq for HeapState {}
|
||||
|
||||
fn dijkstra(
|
||||
edges: &[Vec<(usize, f64)>],
|
||||
start: usize,
|
||||
maybe_goals: &Option<BTreeSet<usize>>,
|
||||
) -> Vec<(usize, f64, Vec<usize>)> {
|
||||
let mut distance = vec![f64::INFINITY; edges.len()];
|
||||
let mut heap = BinaryHeap::new();
|
||||
let mut back_pointers = vec![usize::MAX; edges.len()];
|
||||
distance[start] = 0.;
|
||||
heap.push(HeapState {
|
||||
cost: 0.,
|
||||
node: start,
|
||||
});
|
||||
let mut goals_remaining = maybe_goals.clone();
|
||||
|
||||
while let Some(state) = heap.pop() {
|
||||
if state.cost > distance[state.node] {
|
||||
continue;
|
||||
}
|
||||
|
||||
for (nxt_node, path_weight) in &edges[state.node] {
|
||||
let nxt_cost = state.cost + *path_weight;
|
||||
if nxt_cost < distance[*nxt_node] {
|
||||
heap.push(HeapState {
|
||||
cost: nxt_cost,
|
||||
node: *nxt_node,
|
||||
});
|
||||
distance[*nxt_node] = nxt_cost;
|
||||
back_pointers[*nxt_node] = state.node;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(goals) = &mut goals_remaining {
|
||||
if goals.remove(&state.node) {
|
||||
if goals.is_empty() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let targets = if let Some(goals) = maybe_goals {
|
||||
Left(goals.iter().cloned())
|
||||
} else {
|
||||
Right(0..edges.len())
|
||||
};
|
||||
let ret = targets
|
||||
.map(|target| {
|
||||
let cost = distance[target];
|
||||
if !cost.is_finite() {
|
||||
(target, cost, vec![])
|
||||
} else {
|
||||
let mut path = vec![];
|
||||
let mut current = target;
|
||||
while current != start {
|
||||
path.push(current);
|
||||
current = back_pointers[current];
|
||||
}
|
||||
path.push(start);
|
||||
path.reverse();
|
||||
(target, cost, path)
|
||||
}
|
||||
})
|
||||
.collect_vec();
|
||||
|
||||
ret
|
||||
}
|
||||
Loading…
Reference in New Issue