minimal spanning trees
parent
271387bc59
commit
31549fcb2d
@ -0,0 +1,116 @@
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use itertools::Itertools;
|
||||
use ordered_float::OrderedFloat;
|
||||
use priority_queue::PriorityQueue;
|
||||
use smartstring::{LazyCompact, SmartString};
|
||||
|
||||
use crate::algo::AlgoImpl;
|
||||
use crate::data::expr::Expr;
|
||||
use crate::data::program::{MagicAlgoRuleArg, MagicSymbol};
|
||||
use crate::data::tuple::Tuple;
|
||||
use crate::data::value::DataValue;
|
||||
use crate::runtime::derived::DerivedRelStore;
|
||||
use crate::runtime::transact::SessionTx;
|
||||
|
||||
pub(crate) struct MinimumSpanningTreeKruskal;
|
||||
|
||||
impl AlgoImpl for MinimumSpanningTreeKruskal {
|
||||
fn run(
|
||||
&mut self,
|
||||
tx: &SessionTx,
|
||||
rels: &[MagicAlgoRuleArg],
|
||||
_opts: &BTreeMap<SmartString<LazyCompact>, Expr>,
|
||||
stores: &BTreeMap<MagicSymbol, DerivedRelStore>,
|
||||
out: &DerivedRelStore,
|
||||
) -> Result<()> {
|
||||
let edges = rels
|
||||
.get(0)
|
||||
.ok_or_else(|| anyhow!("'minimum_spanning_tree_kruskal' requires edge relation"))?;
|
||||
let (graph, indices, _, _) =
|
||||
edges.convert_edge_to_weighted_graph(true, true, tx, stores)?;
|
||||
if graph.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let msp = kruskal(&graph);
|
||||
for (src, dst, cost) in msp {
|
||||
out.put(
|
||||
Tuple(vec![
|
||||
indices[src].clone(),
|
||||
indices[dst].clone(),
|
||||
DataValue::from(cost),
|
||||
]),
|
||||
0,
|
||||
);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn kruskal(edges: &[Vec<(usize, f64)>]) -> Vec<(usize, usize, f64)> {
|
||||
let mut pq = PriorityQueue::new();
|
||||
let mut uf = UnionFind::new(edges.len());
|
||||
let mut mst = Vec::with_capacity(edges.len() - 1);
|
||||
for (from, tos) in edges.iter().enumerate() {
|
||||
for (to, cost) in tos {
|
||||
pq.push((from, *to), Reverse(OrderedFloat(*cost)));
|
||||
}
|
||||
}
|
||||
while let Some(((from, to), Reverse(OrderedFloat(cost)))) = pq.pop() {
|
||||
if uf.connected(from, to) {
|
||||
continue;
|
||||
}
|
||||
uf.union(from, to);
|
||||
|
||||
mst.push((from, to, cost));
|
||||
if uf.szs[0] == edges.len() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
mst
|
||||
}
|
||||
|
||||
struct UnionFind {
|
||||
ids: Vec<usize>,
|
||||
szs: Vec<usize>,
|
||||
}
|
||||
|
||||
impl UnionFind {
|
||||
fn new(n: usize) -> Self {
|
||||
Self {
|
||||
ids: (0..n).collect_vec(),
|
||||
szs: vec![1; n],
|
||||
}
|
||||
}
|
||||
fn union(&mut self, p: usize, q: usize) {
|
||||
let root1 = self.find(p);
|
||||
let root2 = self.find(q);
|
||||
if root1 != root2 {
|
||||
if self.szs[root1] < self.szs[root2] {
|
||||
self.szs[root2] += self.szs[root1];
|
||||
self.ids[root1] = root2;
|
||||
} else {
|
||||
self.szs[root1] += self.szs[root2];
|
||||
self.ids[root2] = root1;
|
||||
}
|
||||
}
|
||||
}
|
||||
fn find(&mut self, mut p: usize) -> usize {
|
||||
let mut root = p;
|
||||
while root != self.ids[root] {
|
||||
root = self.ids[root];
|
||||
}
|
||||
while p != root {
|
||||
let next = self.ids[p];
|
||||
self.ids[p] = root;
|
||||
p = next;
|
||||
}
|
||||
root
|
||||
}
|
||||
fn connected(&mut self, p: usize, q: usize) -> bool {
|
||||
self.find(p) == self.find(q)
|
||||
}
|
||||
}
|
@ -0,0 +1,78 @@
|
||||
use std::cmp::Reverse;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use anyhow::{anyhow, Result};
|
||||
use ordered_float::OrderedFloat;
|
||||
use priority_queue::PriorityQueue;
|
||||
use smartstring::{LazyCompact, SmartString};
|
||||
|
||||
use crate::algo::AlgoImpl;
|
||||
use crate::data::expr::Expr;
|
||||
use crate::data::program::{MagicAlgoRuleArg, MagicSymbol};
|
||||
use crate::data::tuple::Tuple;
|
||||
use crate::data::value::DataValue;
|
||||
use crate::runtime::derived::DerivedRelStore;
|
||||
use crate::runtime::transact::SessionTx;
|
||||
|
||||
pub(crate) struct MinimumSpanningTreePrim;
|
||||
|
||||
impl AlgoImpl for MinimumSpanningTreePrim {
|
||||
fn run(
|
||||
&mut self,
|
||||
tx: &SessionTx,
|
||||
rels: &[MagicAlgoRuleArg],
|
||||
_opts: &BTreeMap<SmartString<LazyCompact>, Expr>,
|
||||
stores: &BTreeMap<MagicSymbol, DerivedRelStore>,
|
||||
out: &DerivedRelStore,
|
||||
) -> Result<()> {
|
||||
let edges = rels
|
||||
.get(0)
|
||||
.ok_or_else(|| anyhow!("'minimum_spanning_tree_prim' requires edge relation"))?;
|
||||
let (graph, indices, _, _) =
|
||||
edges.convert_edge_to_weighted_graph(true, true, tx, stores)?;
|
||||
if graph.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let msp = prim(&graph);
|
||||
for (src, dst, cost) in msp {
|
||||
out.put(
|
||||
Tuple(vec![
|
||||
indices[src].clone(),
|
||||
indices[dst].clone(),
|
||||
DataValue::from(cost),
|
||||
]),
|
||||
0,
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn prim(graph: &[Vec<(usize, f64)>]) -> Vec<(usize, usize, f64)> {
|
||||
let mut visited = vec![false; graph.len()];
|
||||
let mut mst_edges = Vec::with_capacity(graph.len() - 1);
|
||||
let mut pq = PriorityQueue::new();
|
||||
|
||||
let mut relax_edges_at_node = |node: usize, pq: &mut PriorityQueue<_, _>| {
|
||||
visited[node] = true;
|
||||
let edges = &graph[node];
|
||||
for (to_node, cost) in edges {
|
||||
if visited[*to_node] {
|
||||
continue;
|
||||
}
|
||||
pq.push_increase(to_node.clone(), (Reverse(OrderedFloat(*cost)), node));
|
||||
}
|
||||
};
|
||||
|
||||
relax_edges_at_node(0, &mut pq);
|
||||
|
||||
while let Some((to_node, (Reverse(OrderedFloat(cost)), from_node))) = pq.pop() {
|
||||
if mst_edges.len() == graph.len() - 1 {
|
||||
break;
|
||||
}
|
||||
mst_edges.push((from_node, to_node, cost));
|
||||
relax_edges_at_node(to_node, &mut pq);
|
||||
}
|
||||
|
||||
mst_edges
|
||||
}
|
Loading…
Reference in New Issue