1164 Advanced

Shortest Path Algorithm with Functional Priority Queue

Graphs

Tutorial

The Problem

Implement Dijkstra's shortest-path algorithm using a functional priority queue — an ordered set that provides O(log n) insert and min-element extraction, mirroring how OCaml uses Set.Make as a persistent priority queue.

🎯 Learning Outcomes

• How Rust's BinaryHeap<Reverse<...>> replaces OCaml's Set.Make-based priority queue

• Why lazy deletion (skipping stale heap entries) is idiomatic in both languages

• How BTreeSet can mirror OCaml's ordered-set priority queue in functional style

• The ownership pattern of cloning to release borrows before mutating collections

🦀 The Rust Way

The idiomatic Rust version uses BinaryHeap<Reverse<(u32, String)>> — a max-heap wrapped with Reverse to invert the ordering into a min-heap. Stale entries (where a shorter path was found later) are skipped lazily when popped. The functional version uses BTreeSet<(u32, String)> which directly mirrors OCaml's Set.Make: a sorted set where the smallest (distance, node) pair is always first in iteration order.

Code Example

pub fn dijkstra(graph: &Graph, start: &str) -> HashMap<String, u32> {
    let mut dist: HashMap<String, u32> = HashMap::new();
    dist.insert(start.to_string(), 0);

    let mut heap: BinaryHeap<Reverse<(u32, String)>> = BinaryHeap::new();
    heap.push(Reverse((0, start.to_string())));

    while let Some(Reverse((d, u))) = heap.pop() {
        if d > dist.get(&u).copied().unwrap_or(u32::MAX) {
            continue;
        }
        if let Some(neighbors) = graph.get(&u) {
            for (v, w) in neighbors {
                let alt = d + w;
                if alt < dist.get(v).copied().unwrap_or(u32::MAX) {
                    dist.insert(v.clone(), alt);
                    heap.push(Reverse((alt, v.clone())));
                }
            }
        }
    }
    dist
}

module PQ = Set.Make(struct
  type t = int * string
  let compare (d1,n1) (d2,n2) = compare (d1,n1) (d2,n2)
end)
module SMap = Map.Make(String)

let dijkstra graph start =
  let dist = SMap.singleton start 0 in
  let pq = PQ.singleton (0, start) in
  let rec go pq dist =
    if PQ.is_empty pq then dist
    else
      let (d, u) = PQ.min_elt pq in
      let pq = PQ.remove (d, u) pq in
      let neighbors = try SMap.find u graph with Not_found -> [] in
      let dist, pq = List.fold_left (fun (dist, pq) (v, w) ->
        let alt = d + w in
        let current = try SMap.find v dist with Not_found -> max_int in
        if alt < current then
          (SMap.add v alt dist, PQ.add (alt, v) pq)
        else (dist, pq)
      ) (dist, pq) neighbors in
      go pq dist
  in go pq dist

Key Differences

Priority queue representation: OCaml uses Set.Make (persistent ordered set); Rust uses BinaryHeap<Reverse<...>> (imperative) or BTreeSet (structural analog to Set.Make)

Mutability: OCaml threads immutable maps through recursion; Rust mutates HashMap/BTreeMap and the heap in place

Stale entry handling: Both use lazy deletion — OCaml checks if d > dist_u then go pq' dist, Rust checks if d > dist[u] { continue }

Borrow constraint: Rust must .cloned() the BTreeSet's first element to release the immutable borrow before calling .remove() — OCaml's persistent data structures have no such constraint

OCaml Approach

OCaml uses Set.Make with a (int * string) comparison to create a sorted set that acts as a functional priority queue — always ordered by (distance, node), with the minimum element retrievable in O(log n). The algorithm is tail-recursive via a local go function, threading the priority queue and distance map as immutable values that are replaced on each step.

Full Source

#![allow(clippy::all)]
use std::cmp::Reverse;
use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap};

/// Graph as adjacency list: node name → list of (neighbor, weight)
pub type Graph = HashMap<String, Vec<(String, u32)>>;

/// Dijkstra's shortest-path — idiomatic Rust
///
/// Uses `BinaryHeap<Reverse<...>>` as a standard min-heap.
/// Stale entries are discarded lazily when popped (lazy deletion).
pub fn dijkstra(graph: &Graph, start: &str) -> HashMap<String, u32> {
    let mut dist: HashMap<String, u32> = HashMap::new();
    dist.insert(start.to_string(), 0);

    // Reverse turns the max-heap into a min-heap: smallest distance pops first
    let mut heap: BinaryHeap<Reverse<(u32, String)>> = BinaryHeap::new();
    heap.push(Reverse((0, start.to_string())));

    while let Some(Reverse((d, u))) = heap.pop() {
        // Lazy deletion: skip stale entries where a shorter path was already recorded
        if d > dist.get(&u).copied().unwrap_or(u32::MAX) {
            continue;
        }

        if let Some(neighbors) = graph.get(&u) {
            for (v, w) in neighbors {
                let alt = d + w;
                if alt < dist.get(v).copied().unwrap_or(u32::MAX) {
                    dist.insert(v.clone(), alt);
                    heap.push(Reverse((alt, v.clone())));
                }
            }
        }
    }

    dist
}

/// Graph for the functional implementation (BTreeMap for sorted, deterministic output)
pub type FunctionalGraph = BTreeMap<String, Vec<(String, u32)>>;

/// Dijkstra's shortest-path — functional style mirroring OCaml's Set.Make PQ
///
/// Uses `BTreeSet<(u32, String)>` as an ordered set, exactly as OCaml uses
/// `Set.Make` for the priority queue. The smallest element (min_elt) is always
/// the first element in BTreeSet's iteration order.
pub fn dijkstra_functional(graph: &FunctionalGraph, start: &str) -> BTreeMap<String, u32> {
    let mut dist: BTreeMap<String, u32> = BTreeMap::new();
    dist.insert(start.to_string(), 0);

    // BTreeSet mimics OCaml's Set.Make — ordered set, min element is first
    let mut pq: BTreeSet<(u32, String)> = BTreeSet::new();
    pq.insert((0, start.to_string()));

    // OCaml: if PQ.is_empty pq then dist else ...
    // .cloned() releases the immutable borrow on pq so we can mutate it
    while let Some(entry) = pq.iter().next().cloned() {
        // OCaml: PQ.remove (d, u) pq
        pq.remove(&entry);
        let (d, u) = entry;

        // OCaml: if d > dist_u then go pq' dist (skip stale entries)
        if d > dist.get(&u).copied().unwrap_or(u32::MAX) {
            continue;
        }

        // OCaml: List.fold_left over neighbors updating dist and pq
        if let Some(neighbors) = graph.get(&u) {
            for (v, w) in neighbors {
                let alt = d + w;
                if alt < dist.get(v).copied().unwrap_or(u32::MAX) {
                    dist.insert(v.clone(), alt);
                    pq.insert((alt, v.clone()));
                }
            }
        }
    }

    dist
}

#[cfg(test)]
mod tests {
    use super::*;

    fn sample_graph() -> Graph {
        let mut g: Graph = HashMap::new();
        g.insert("a".into(), vec![("b".into(), 1), ("c".into(), 4)]);
        g.insert("b".into(), vec![("c".into(), 2), ("d".into(), 6)]);
        g.insert("c".into(), vec![("d".into(), 3)]);
        g.insert("d".into(), vec![]);
        g
    }

    fn sample_functional_graph() -> FunctionalGraph {
        let mut g: FunctionalGraph = BTreeMap::new();
        g.insert("a".into(), vec![("b".into(), 1), ("c".into(), 4)]);
        g.insert("b".into(), vec![("c".into(), 2), ("d".into(), 6)]);
        g.insert("c".into(), vec![("d".into(), 3)]);
        g.insert("d".into(), vec![]);
        g
    }

    #[test]
    fn test_start_node_has_zero_distance() {
        let dist = dijkstra(&sample_graph(), "a");
        assert_eq!(dist.get("a").copied(), Some(0));
    }

    #[test]
    fn test_shortest_paths_indirect_beats_direct() {
        // a->c direct = 4, but a->b->c = 1+2 = 3 wins
        let dist = dijkstra(&sample_graph(), "a");
        assert_eq!(dist.get("a").copied(), Some(0));
        assert_eq!(dist.get("b").copied(), Some(1)); // a->b
        assert_eq!(dist.get("c").copied(), Some(3)); // a->b->c, not a->c=4
        assert_eq!(dist.get("d").copied(), Some(6)); // a->b->c->d
    }

    #[test]
    fn test_unreachable_node_absent_from_result() {
        let mut g = sample_graph();
        g.insert("z".into(), vec![]); // z is disconnected
        let dist = dijkstra(&g, "a");
        assert_eq!(dist.get("z"), None);
    }

    #[test]
    fn test_single_node_graph() {
        let mut g: Graph = HashMap::new();
        g.insert("x".into(), vec![]);
        let dist = dijkstra(&g, "x");
        assert_eq!(dist.get("x").copied(), Some(0));
        assert_eq!(dist.len(), 1);
    }

    #[test]
    fn test_functional_matches_idiomatic() {
        let dist_std = dijkstra(&sample_graph(), "a");
        let dist_btree = dijkstra_functional(&sample_functional_graph(), "a");
        for key in ["a", "b", "c", "d"] {
            assert_eq!(
                dist_std.get(key),
                dist_btree.get(key),
                "mismatch for node {key}"
            );
        }
    }

    #[test]
    fn test_functional_start_zero_unreachable_absent() {
        let mut g: FunctionalGraph = BTreeMap::new();
        g.insert("a".into(), vec![("b".into(), 5)]);
        g.insert("b".into(), vec![]);
        g.insert("c".into(), vec![]); // disconnected
        let dist = dijkstra_functional(&g, "a");
        assert_eq!(dist.get("a").copied(), Some(0));
        assert_eq!(dist.get("b").copied(), Some(5));
        assert_eq!(dist.get("c"), None);
    }
}

module PQ = Set.Make(struct
  type t = int * string
  let compare (d1,n1) (d2,n2) = compare (d1,n1) (d2,n2)
end)
module SMap = Map.Make(String)

let dijkstra graph start =
  let dist = SMap.singleton start 0 in
  let pq = PQ.singleton (0, start) in
  let rec go pq dist =
    if PQ.is_empty pq then dist
    else
      let (d, u) = PQ.min_elt pq in
      let pq = PQ.remove (d, u) pq in
      let neighbors = try SMap.find u graph with Not_found -> [] in
      let dist, pq = List.fold_left (fun (dist, pq) (v, w) ->
        let alt = d + w in
        let current = try SMap.find v dist with Not_found -> max_int in
        if alt < current then
          (SMap.add v alt dist, PQ.add (alt, v) pq)
        else (dist, pq)
      ) (dist, pq) neighbors in
      go pq dist
  in go pq dist

let () =
  let g = SMap.of_list [
    ("a",[("b",1);("c",4)]); ("b",[("c",2);("d",6)]);
    ("c",[("d",3)]); ("d",[])
  ] in
  let dist = dijkstra g "a" in
  SMap.iter (Printf.printf "%s: %d\n") dist

✓ Tests Rust test suite

#[cfg(test)]
mod tests {
    use super::*;

    fn sample_graph() -> Graph {
        let mut g: Graph = HashMap::new();
        g.insert("a".into(), vec![("b".into(), 1), ("c".into(), 4)]);
        g.insert("b".into(), vec![("c".into(), 2), ("d".into(), 6)]);
        g.insert("c".into(), vec![("d".into(), 3)]);
        g.insert("d".into(), vec![]);
        g
    }

    fn sample_functional_graph() -> FunctionalGraph {
        let mut g: FunctionalGraph = BTreeMap::new();
        g.insert("a".into(), vec![("b".into(), 1), ("c".into(), 4)]);
        g.insert("b".into(), vec![("c".into(), 2), ("d".into(), 6)]);
        g.insert("c".into(), vec![("d".into(), 3)]);
        g.insert("d".into(), vec![]);
        g
    }

    #[test]
    fn test_start_node_has_zero_distance() {
        let dist = dijkstra(&sample_graph(), "a");
        assert_eq!(dist.get("a").copied(), Some(0));
    }

    #[test]
    fn test_shortest_paths_indirect_beats_direct() {
        // a->c direct = 4, but a->b->c = 1+2 = 3 wins
        let dist = dijkstra(&sample_graph(), "a");
        assert_eq!(dist.get("a").copied(), Some(0));
        assert_eq!(dist.get("b").copied(), Some(1)); // a->b
        assert_eq!(dist.get("c").copied(), Some(3)); // a->b->c, not a->c=4
        assert_eq!(dist.get("d").copied(), Some(6)); // a->b->c->d
    }

    #[test]
    fn test_unreachable_node_absent_from_result() {
        let mut g = sample_graph();
        g.insert("z".into(), vec![]); // z is disconnected
        let dist = dijkstra(&g, "a");
        assert_eq!(dist.get("z"), None);
    }

    #[test]
    fn test_single_node_graph() {
        let mut g: Graph = HashMap::new();
        g.insert("x".into(), vec![]);
        let dist = dijkstra(&g, "x");
        assert_eq!(dist.get("x").copied(), Some(0));
        assert_eq!(dist.len(), 1);
    }

    #[test]
    fn test_functional_matches_idiomatic() {
        let dist_std = dijkstra(&sample_graph(), "a");
        let dist_btree = dijkstra_functional(&sample_functional_graph(), "a");
        for key in ["a", "b", "c", "d"] {
            assert_eq!(
                dist_std.get(key),
                dist_btree.get(key),
                "mismatch for node {key}"
            );
        }
    }

    #[test]
    fn test_functional_start_zero_unreachable_absent() {
        let mut g: FunctionalGraph = BTreeMap::new();
        g.insert("a".into(), vec![("b".into(), 5)]);
        g.insert("b".into(), vec![]);
        g.insert("c".into(), vec![]); // disconnected
        let dist = dijkstra_functional(&g, "a");
        assert_eq!(dist.get("a").copied(), Some(0));
        assert_eq!(dist.get("b").copied(), Some(5));
        assert_eq!(dist.get("c"), None);
    }
}

Deep Comparison

OCaml vs Rust: Shortest Path with Functional Priority Queue

Side-by-Side Code

OCaml

module PQ = Set.Make(struct
  type t = int * string
  let compare (d1,n1) (d2,n2) = compare (d1,n1) (d2,n2)
end)
module SMap = Map.Make(String)

let dijkstra graph start =
  let dist = SMap.singleton start 0 in
  let pq = PQ.singleton (0, start) in
  let rec go pq dist =
    if PQ.is_empty pq then dist
    else
      let (d, u) = PQ.min_elt pq in
      let pq = PQ.remove (d, u) pq in
      let neighbors = try SMap.find u graph with Not_found -> [] in
      let dist, pq = List.fold_left (fun (dist, pq) (v, w) ->
        let alt = d + w in
        let current = try SMap.find v dist with Not_found -> max_int in
        if alt < current then
          (SMap.add v alt dist, PQ.add (alt, v) pq)
        else (dist, pq)
      ) (dist, pq) neighbors in
      go pq dist
  in go pq dist

Rust (idiomatic — BinaryHeap min-heap)

pub fn dijkstra(graph: &Graph, start: &str) -> HashMap<String, u32> {
    let mut dist: HashMap<String, u32> = HashMap::new();
    dist.insert(start.to_string(), 0);

    let mut heap: BinaryHeap<Reverse<(u32, String)>> = BinaryHeap::new();
    heap.push(Reverse((0, start.to_string())));

    while let Some(Reverse((d, u))) = heap.pop() {
        if d > dist.get(&u).copied().unwrap_or(u32::MAX) {
            continue;
        }
        if let Some(neighbors) = graph.get(&u) {
            for (v, w) in neighbors {
                let alt = d + w;
                if alt < dist.get(v).copied().unwrap_or(u32::MAX) {
                    dist.insert(v.clone(), alt);
                    heap.push(Reverse((alt, v.clone())));
                }
            }
        }
    }
    dist
}

Rust (functional — BTreeSet mirrors OCaml's Set.Make)

pub fn dijkstra_functional(graph: &FunctionalGraph, start: &str) -> BTreeMap<String, u32> {
    let mut dist: BTreeMap<String, u32> = BTreeMap::new();
    dist.insert(start.to_string(), 0);

    let mut pq: BTreeSet<(u32, String)> = BTreeSet::new();
    pq.insert((0, start.to_string()));

    while let Some(entry) = pq.iter().next().cloned() {
        pq.remove(&entry);
        let (d, u) = entry;

        if d > dist.get(&u).copied().unwrap_or(u32::MAX) {
            continue;
        }
        if let Some(neighbors) = graph.get(&u) {
            for (v, w) in neighbors {
                let alt = d + w;
                if alt < dist.get(v).copied().unwrap_or(u32::MAX) {
                    dist.insert(v.clone(), alt);
                    pq.insert((alt, v.clone()));
                }
            }
        }
    }
    dist
}

Type Signatures

Concept	OCaml	Rust (idiomatic)
Graph type	`(string * (string * int) list) SMap.t`	`HashMap<String, Vec<(String, u32)>>`
Distance map	`int SMap.t`	`HashMap<String, u32>`
Priority queue	`PQ.t` (`Set.Make` of `int * string`)	`BinaryHeap<Reverse<(u32, String)>>`
Functional PQ	`PQ.t` (same — always functional)	`BTreeSet<(u32, String)>`
Function signature	`dijkstra : ... SMap.t -> string -> int SMap.t`	`fn dijkstra(graph: &Graph, start: &str) -> HashMap<String, u32>`

Key Insights

**OCaml's Set.Make is a persistent sorted set used as a priority queue.** It keeps elements in sorted order by the custom comparator (distance, node), so PQ.min_elt is O(log n). Rust's BTreeSet<(u32, String)> is the direct structural analog — BTreeSet is also a sorted ordered set, and its first element in iteration order is always the minimum.

**BinaryHeap<Reverse<...>> is the idiomatic Rust choice.** It's a standard binary heap with O(log n) push and pop. Wrapping with Reverse inverts the max-heap ordering into a min-heap. This is more cache-friendly and lower-overhead than BTreeSet for the priority queue use case.

Lazy deletion is identical in both languages. OCaml checks if d > dist_u then go pq' dist to skip stale PQ entries. Rust checks if d > dist[u] { continue }. Neither eagerly removes outdated entries from the queue — they're just discarded when encountered.

Borrow checker changes the BTreeSet access pattern. In OCaml, PQ.min_elt pq and PQ.remove (d,u) pq are two separate operations on an immutable structure — no aliasing concern. In Rust, pq.iter().next() borrows pq immutably; we must call .cloned() to obtain an owned copy and release the borrow before calling pq.remove(...). This is a direct consequence of Rust's aliasing rules.

Immutable threading vs. mutable update. OCaml's go pq dist passes updated pq and dist as new values in each recursive call — purely functional. Rust mutates dist and heap/pq in place inside a while loop. The observable behavior is identical; the implementation strategy differs completely.

When to Use Each Style

Use idiomatic Rust (BinaryHeap) when: you want maximum performance for graph algorithms — BinaryHeap is O(log n) push/pop with good cache behavior, and is what production Rust graph libraries use.

Use functional Rust (BTreeSet) when: you want the code structure to closely mirror an OCaml or Haskell implementation for pedagogical purposes, or when you need the set semantics (e.g., guaranteed uniqueness of entries, range queries) beyond just priority ordering.

Exercises

Replace the functional priority queue with a BinaryHeap-based one and verify the results are identical while measuring the performance difference.

Implement the k-shortest paths algorithm (Yen's algorithm) on top of this Dijkstra implementation, returning the k shortest simple paths from source to target.

Model a transportation network (cities as nodes, roads as edges with km distances) and use the shortest-path algorithm to find optimal routes, then extend to support waypoints (must-visit intermediate nodes).

Open Source Repos

functional-rust

View the source for this example on GitHub — OCaml and Rust side by side in the repo.

Rust