minimal spanning trees

README.md

@@ -6,11 +6,11 @@
   * [x] bfs
   * [x] dfs
   * [x] shortest path
-  * [ ] A*
+  * [x] A*
   * [x] Yen's k-shortest
   * [ ] all-pairs shortest path
   * [x] single-source shortest path
-  * [ ] minimum spanning tree
+  * [x] minimum spanning tree
   * [ ] random walking
   * [x] degree centrality
   * [ ] closeness centrality

cozo_webui/src/App.svelte

@@ -25,13 +25,12 @@
                     body: query
                 });
 
-                let time = Math.round(performance.now() - started);
-                statusMessage = `finished in ${time}ms`
-
                 if (!response.ok) {
                     throw await response.text();
                 }
                 let res = await response.json();
+                statusMessage = `finished in ${res.time_taken}ms`
+
                 if (res.rows) {
                     if (!res.headers) {
                         res.headers = [];

cozohttp/src/main.rs

@@ -5,12 +5,13 @@ use std::time::Instant;
 use actix_cors::Cors;
 use actix_web::rt::task::spawn_blocking;
 use actix_web::{post, web, App, HttpResponse, HttpServer, Responder};
+use actix_web_static_files::ResourceFiles;
+use anyhow::anyhow;
 use clap::Parser;
 use env_logger::Env;
 use log::info;
+use serde_json::json;
 
-use actix_web_static_files::ResourceFiles;
-use anyhow::anyhow;
 use cozo::{Db, DbBuilder};
 
 type Result<T> = std::result::Result<T, RespError>;
@@ -71,8 +72,10 @@ async fn query(body: web::Bytes, data: web::Data<AppStateWithDb>) -> Result<impl
     let db = data.db.new_session()?;
     let start = Instant::now();
     let task = spawn_blocking(move || db.run_script(&text));
-    let result = task.await.map_err(|e| anyhow!(e))??;
-    info!("finished query in {:?}", start.elapsed());
+    let mut result = task.await.map_err(|e| anyhow!(e))??;
+    if let Some(obj) = result.as_object_mut() {
+        obj.insert("time_taken".to_string(), json!(start.elapsed().as_millis() as u64));
+    }
     Ok(HttpResponse::Ok().json(result))
 }
 

src/algo/kruskal.rs

@@ -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)
+    }
+}

src/algo/mod.rs

@@ -8,6 +8,8 @@ use crate::algo::astar::ShortestPathAStar;
 use crate::algo::bfs::Bfs;
 use crate::algo::degree_centrality::DegreeCentrality;
 use crate::algo::dfs::Dfs;
+use crate::algo::kruskal::MinimumSpanningTreeKruskal;
+use crate::algo::prim::MinimumSpanningTreePrim;
 use crate::algo::shortest_path_dijkstra::ShortestPathDijkstra;
 use crate::algo::strongly_connected_components::StronglyConnectedComponent;
 use crate::algo::top_sort::TopSort;
@@ -25,7 +27,9 @@ pub(crate) mod astar;
 pub(crate) mod bfs;
 pub(crate) mod degree_centrality;
 pub(crate) mod dfs;
+pub(crate) mod kruskal;
 pub(crate) mod page_rank;
+pub(crate) mod prim;
 pub(crate) mod shortest_path_dijkstra;
 pub(crate) mod strongly_connected_components;
 pub(crate) mod top_sort;
@@ -62,6 +66,8 @@ impl AlgoHandle {
             "shortest_path_dijkstra" => 4,
             "shortest_path_astar" => 4,
             "k_shortest_path_yen" => 4,
+            "minimum_spanning_tree_prim" => 3,
+            "minimum_spanning_tree_kruskal" => 3,
             "top_sort" => 2,
             "connected_components" => 2,
             "strongly_connected_components" | "scc" => 2,
@@ -78,6 +84,8 @@ impl AlgoHandle {
             "shortest_path_dijkstra" => Box::new(ShortestPathDijkstra),
             "shortest_path_astar" => Box::new(ShortestPathAStar),
             "k_shortest_path_yen" => Box::new(KShortestPathYen),
+            "minimum_spanning_tree_prim" => Box::new(MinimumSpanningTreePrim),
+            "minimum_spanning_tree_kruskal" => Box::new(MinimumSpanningTreeKruskal),
             "top_sort" => Box::new(TopSort),
             "connected_components" => Box::new(StronglyConnectedComponent::new(false)),
             "strongly_connected_components" | "scc" => {

src/algo/prim.rs

@@ -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
+}