1169 Advanced

Dijkstra's Shortest Path with a Priority Queue

GraphsAlgorithms

Tutorial

The Problem

Find the shortest distances from a start node to all reachable nodes in a weighted directed graph, using a priority-queue–driven greedy relaxation loop.

🎯 Learning Outcomes

• How OCaml's Set.Make ordered set (used as a min-priority-queue) maps to

Rust's BinaryHeap<Reverse<…>> min-heap.

• How OCaml's Map.Make(String) immutable map maps to Rust's HashMap.

• Why Rust needs the Reverse wrapper: BinaryHeap is a max-heap by default.

• How OCaml's tail-recursive rec go pq dist pattern translates to Rust's

imperative while let … = heap.pop() loop (idiomatic) or a recursive helper that passes accumulator state forward (functional style).

🦀 The Rust Way

Rust's BinaryHeap is a max-heap; wrapping each entry in std::cmp::Reverse flips the ordering so pop() yields the smallest distance first — exactly what Set.min_elt does in OCaml. HashMap is mutable and updated in place. The idiomatic Rust version uses a while let Some(…) = heap.pop() loop; the functional version passes (heap, dist) through a recursive helper, mirroring OCaml's rec go directly and using .iter().fold(…) for the neighbour relaxation that OCaml expresses with List.fold_left.

Code Example

pub fn dijkstra<'a>(graph: &Graph<'a>, start: &'a str) -> HashMap<&'a str, u32> {
    let mut dist: HashMap<&str, u32> = HashMap::new();
    dist.insert(start, 0);
    let mut heap: BinaryHeap<Reverse<(u32, &str)>> = BinaryHeap::new();
    heap.push(Reverse((0, start)));

    while let Some(Reverse((d, u))) = heap.pop() {
        if d > *dist.get(u).unwrap_or(&u32::MAX) {
            continue;                    // stale entry — skip
        }
        let neighbors = graph.get(u).map(Vec::as_slice).unwrap_or(&[]);
        for &(v, w) in neighbors {
            let alt = d + w;
            if alt < *dist.get(v).unwrap_or(&u32::MAX) {
                dist.insert(v, alt);
                heap.push(Reverse((alt, v)));
            }
        }
    }
    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: OCaml uses Set.Make (balanced BST, O(log n) insert/delete-min); Rust uses BinaryHeap (binary heap, O(log n) push/pop) — same asymptotic cost, different data structure.

Min-heap polarity: OCaml's Set.min_elt naturally returns the smallest element; Rust's BinaryHeap is max-first, requiring Reverse to achieve the same behaviour.

Stale entries: OCaml's Set.remove ensures each (distance, node) pair is unique, preventing re-processing. Rust's heap allows duplicates, so the idiomatic version guards with if d > dist[u] { continue }.

Mutation vs immutability: OCaml's SMap.add and PQ.add return new values; Rust's HashMap::insert and BinaryHeap::push mutate in place — the result is the same, but the Rust version avoids allocation per update.

OCaml Approach

OCaml uses two functorised modules: Set.Make ordered by (distance, node) acts as a self-sorting priority queue (min element accessible with min_elt), and Map.Make(String) stores distances immutably. The algorithm is a tail-recursive go function that shadows pq and dist on every step, so the whole loop is expressed as pure data transformation.

Full Source

#![allow(clippy::all)]
//! # Dijkstra's Shortest Path with a Priority Queue
//!
//! OCaml uses `Set.Make` as a sorted priority queue (ordered by `(dist, node)`
//! tuple) and `Map.Make(String)` as an immutable distance map, threading both
//! through a tail-recursive `go` loop.
//!
//! Rust's `std::collections::BinaryHeap` is a max-heap; wrapping entries in
//! `std::cmp::Reverse` turns it into the min-heap Dijkstra requires.
//! `HashMap` plays the role of OCaml's `SMap`.
//!
//! This module shows:
//! 1. Idiomatic Rust — imperative `while let` loop, mutable `HashMap`
//! 2. Functional style — immutable-ish accumulator passed through a helper
//! 3. The OCaml→Rust translation of `Set.Make` → `BinaryHeap<Reverse<…>>`

use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};

/// Weighted directed graph: node name → list of (neighbour, weight).
///
/// Uses `&str` keys so callers can pass string literals directly.
/// `u32` weights match OCaml's `int` (non-negative distances).
pub type Graph<'a> = HashMap<&'a str, Vec<(&'a str, u32)>>;

// ---------------------------------------------------------------------------
// Solution 1: Idiomatic Rust — `while let` loop with mutable state
//
// OCaml equivalent uses a tail-recursive `go` function that shadows
// `pq` and `dist` on every call. Rust's `while let` over `heap.pop()`
// expresses the same control flow without recursion.
// ---------------------------------------------------------------------------

/// Run Dijkstra from `start` and return shortest distances to all reachable nodes.
///
/// OCaml type: `dijkstra : int SMap.t SMap.t -> string -> int SMap.t`
/// Rust type:  `fn dijkstra<'a>(graph: &Graph<'a>, start: &'a str) -> HashMap<&'a str, u32>`
///
/// Uses `BinaryHeap<Reverse<(u32, &str)>>` as a min-heap.
/// `Reverse` flips the natural max-heap ordering so the smallest distance
/// is always at the top — the same invariant OCaml's `Set.min_elt` provides.
pub fn dijkstra<'a>(graph: &Graph<'a>, start: &'a str) -> HashMap<&'a str, u32> {
    // dist[node] = best known distance from start; unvisited nodes are absent.
    let mut dist: HashMap<&str, u32> = HashMap::new();
    dist.insert(start, 0);

    // Min-heap: Reverse wraps (distance, node) so smaller distances pop first.
    let mut heap: BinaryHeap<Reverse<(u32, &str)>> = BinaryHeap::new();
    heap.push(Reverse((0, start)));

    while let Some(Reverse((d, u))) = heap.pop() {
        // Skip stale entries: a shorter path was already found and processed.
        if d > *dist.get(u).unwrap_or(&u32::MAX) {
            continue;
        }

        let neighbors = graph.get(u).map(Vec::as_slice).unwrap_or(&[]);
        for &(v, w) in neighbors {
            let alt = d + w;
            // Only update if we found a strictly shorter path.
            let current = *dist.get(v).unwrap_or(&u32::MAX);
            if alt < current {
                dist.insert(v, alt);
                heap.push(Reverse((alt, v)));
            }
        }
    }

    dist
}

// ---------------------------------------------------------------------------
// Solution 2: Functional style — immutable accumulator, helper function
//
// Mirrors the OCaml `rec go pq dist` pattern directly.
// Rust doesn't optimise tail calls, so we use an explicit stack (the heap)
// instead of mutual recursion to avoid stack overflow on large graphs.
// The key OCaml idiom `List.fold_left` maps to `.iter().fold(…)`.
// ---------------------------------------------------------------------------

/// Functional-style Dijkstra: state is accumulated and passed forward,
/// mirroring OCaml's `rec go pq dist` tail-recursive helper.
///
/// `fold_neighbors` plays the role of `List.fold_left` in the OCaml source.
pub fn dijkstra_functional<'a>(graph: &Graph<'a>, start: &'a str) -> HashMap<&'a str, u32> {
    let dist: HashMap<&str, u32> = std::iter::once((start, 0u32)).collect();
    let heap: BinaryHeap<Reverse<(u32, &str)>> = std::iter::once(Reverse((0u32, start))).collect();
    go(graph, heap, dist)
}

/// Recursive helper — accumulates `(heap, dist)` toward the fixed point
/// where the heap is empty (all reachable nodes settled).
fn go<'a>(
    graph: &Graph<'a>,
    mut heap: BinaryHeap<Reverse<(u32, &'a str)>>,
    dist: HashMap<&'a str, u32>,
) -> HashMap<&'a str, u32> {
    // Base case: OCaml's `if PQ.is_empty pq then dist`
    let Some(Reverse((d, u))) = heap.pop() else {
        return dist;
    };

    // Stale-entry guard (OCaml's Set avoids duplicates; heap does not)
    if d > *dist.get(u).unwrap_or(&u32::MAX) {
        return go(graph, heap, dist);
    }

    // `List.fold_left` over neighbours → `.iter().fold(…)` over the slice
    let neighbors = graph.get(u).map(Vec::as_slice).unwrap_or(&[]);
    let (dist, heap) = neighbors
        .iter()
        .fold((dist, heap), |(mut d_map, mut h), &(v, w)| {
            let alt = d + w;
            let current = *d_map.get(v).unwrap_or(&u32::MAX);
            if alt < current {
                d_map.insert(v, alt);
                h.push(Reverse((alt, v)));
            }
            (d_map, h)
        });

    go(graph, heap, dist)
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

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

    #[test]
    fn test_single_node_start() {
        // A graph with only the start node reachable.
        let mut g = Graph::new();
        g.insert("x", vec![]);
        let dist = dijkstra(&g, "x");
        assert_eq!(dist["x"], 0);
        assert_eq!(dist.len(), 1);
    }

    #[test]
    fn test_direct_edge() {
        let mut g = Graph::new();
        g.insert("a", vec![("b", 7)]);
        g.insert("b", vec![]);
        let dist = dijkstra(&g, "a");
        assert_eq!(dist["a"], 0);
        assert_eq!(dist["b"], 7);
    }

    #[test]
    fn test_shortest_path_prefers_indirect_route() {
        // a→c directly costs 4, but a→b→c costs 1+2=3 (shorter).
        let g = sample_graph();
        let dist = dijkstra(&g, "a");
        assert_eq!(dist["a"], 0);
        assert_eq!(dist["b"], 1);
        assert_eq!(dist["c"], 3); // via b, not directly
        assert_eq!(dist["d"], 6); // a→b(1)→c(2)→d(3) = 6, not a→b(1)→d(6)=7
    }

    #[test]
    fn test_unreachable_node_absent() {
        // Node "z" is in the graph but unreachable from "a".
        let mut g = sample_graph();
        g.insert("z", vec![("a", 1)]); // z can reach a, but a cannot reach z
        let dist = dijkstra(&g, "a");
        assert!(
            !dist.contains_key("z"),
            "unreachable node must not appear in dist"
        );
    }

    #[test]
    fn test_both_implementations_agree() {
        let g = sample_graph();
        let d1 = dijkstra(&g, "a");
        let d2 = dijkstra_functional(&g, "a");
        assert_eq!(
            d1, d2,
            "idiomatic and functional implementations must agree"
        );
    }

    #[test]
    fn test_disconnected_start_not_in_graph() {
        // start node has no entry in graph — dist should still contain start=0.
        let g: Graph = HashMap::new();
        let dist = dijkstra(&g, "alone");
        assert_eq!(dist["alone"], 0);
        assert_eq!(dist.len(), 1);
    }
}

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<'a>() -> Graph<'a> {
        let mut g = Graph::new();
        g.insert("a", vec![("b", 1), ("c", 4)]);
        g.insert("b", vec![("c", 2), ("d", 6)]);
        g.insert("c", vec![("d", 3)]);
        g.insert("d", vec![]);
        g
    }

    #[test]
    fn test_single_node_start() {
        // A graph with only the start node reachable.
        let mut g = Graph::new();
        g.insert("x", vec![]);
        let dist = dijkstra(&g, "x");
        assert_eq!(dist["x"], 0);
        assert_eq!(dist.len(), 1);
    }

    #[test]
    fn test_direct_edge() {
        let mut g = Graph::new();
        g.insert("a", vec![("b", 7)]);
        g.insert("b", vec![]);
        let dist = dijkstra(&g, "a");
        assert_eq!(dist["a"], 0);
        assert_eq!(dist["b"], 7);
    }

    #[test]
    fn test_shortest_path_prefers_indirect_route() {
        // a→c directly costs 4, but a→b→c costs 1+2=3 (shorter).
        let g = sample_graph();
        let dist = dijkstra(&g, "a");
        assert_eq!(dist["a"], 0);
        assert_eq!(dist["b"], 1);
        assert_eq!(dist["c"], 3); // via b, not directly
        assert_eq!(dist["d"], 6); // a→b(1)→c(2)→d(3) = 6, not a→b(1)→d(6)=7
    }

    #[test]
    fn test_unreachable_node_absent() {
        // Node "z" is in the graph but unreachable from "a".
        let mut g = sample_graph();
        g.insert("z", vec![("a", 1)]); // z can reach a, but a cannot reach z
        let dist = dijkstra(&g, "a");
        assert!(
            !dist.contains_key("z"),
            "unreachable node must not appear in dist"
        );
    }

    #[test]
    fn test_both_implementations_agree() {
        let g = sample_graph();
        let d1 = dijkstra(&g, "a");
        let d2 = dijkstra_functional(&g, "a");
        assert_eq!(
            d1, d2,
            "idiomatic and functional implementations must agree"
        );
    }

    #[test]
    fn test_disconnected_start_not_in_graph() {
        // start node has no entry in graph — dist should still contain start=0.
        let g: Graph = HashMap::new();
        let dist = dijkstra(&g, "alone");
        assert_eq!(dist["alone"], 0);
        assert_eq!(dist.len(), 1);
    }
}

Deep Comparison

OCaml vs Rust: Dijkstra's Shortest Path with a 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)

pub fn dijkstra<'a>(graph: &Graph<'a>, start: &'a str) -> HashMap<&'a str, u32> {
    let mut dist: HashMap<&str, u32> = HashMap::new();
    dist.insert(start, 0);
    let mut heap: BinaryHeap<Reverse<(u32, &str)>> = BinaryHeap::new();
    heap.push(Reverse((0, start)));

    while let Some(Reverse((d, u))) = heap.pop() {
        if d > *dist.get(u).unwrap_or(&u32::MAX) {
            continue;                    // stale entry — skip
        }
        let neighbors = graph.get(u).map(Vec::as_slice).unwrap_or(&[]);
        for &(v, w) in neighbors {
            let alt = d + w;
            if alt < *dist.get(v).unwrap_or(&u32::MAX) {
                dist.insert(v, alt);
                heap.push(Reverse((alt, v)));
            }
        }
    }
    dist
}

Rust (functional/recursive — mirrors OCaml's `rec go`)

pub fn dijkstra_functional<'a>(graph: &Graph<'a>, start: &'a str) -> HashMap<&'a str, u32> {
    let dist = std::iter::once((start, 0u32)).collect();
    let heap = std::iter::once(Reverse((0u32, start))).collect();
    go(graph, heap, dist)
}

fn go<'a>(
    graph: &Graph<'a>,
    mut heap: BinaryHeap<Reverse<(u32, &'a str)>>,
    dist: HashMap<&'a str, u32>,
) -> HashMap<&'a str, u32> {
    let Some(Reverse((d, u))) = heap.pop() else {
        return dist;                     // base case: PQ.is_empty pq
    };
    if d > *dist.get(u).unwrap_or(&u32::MAX) {
        return go(graph, heap, dist);    // stale — tail-recurse unchanged
    }
    let neighbors = graph.get(u).map(Vec::as_slice).unwrap_or(&[]);
    let (dist, heap) = neighbors.iter()  // List.fold_left → .iter().fold(…)
        .fold((dist, heap), |(mut d_map, mut h), &(v, w)| {
            let alt = d + w;
            if alt < *d_map.get(v).unwrap_or(&u32::MAX) {
                d_map.insert(v, alt);
                h.push(Reverse((alt, v)));
            }
            (d_map, h)
        });
    go(graph, heap, dist)
}

Type Signatures

Concept	OCaml	Rust
Function signature	`val dijkstra : int list SMap.t -> string -> int SMap.t`	`fn dijkstra<'a>(graph: &Graph<'a>, start: &'a str) -> HashMap<&'a str, u32>`
Priority queue type	`PQ.t` (`Set.Make` over `(int * string)`)	`BinaryHeap<Reverse<(u32, &str)>>`
Distance map type	`int SMap.t` (`Map.Make(String)`)	`HashMap<&str, u32>`
Sentinel for "infinity"	`max_int`	`u32::MAX` (or absent from map)
Min-element extraction	`PQ.min_elt pq` + `PQ.remove`	`heap.pop()` (single O(log n) step)
Neighbour iteration	`List.fold_left (fun acc (v,w) -> …) init neighbors`	`.iter().fold(init, \\|acc, &(v,w)\\| …)`

Key Insights

**Set.Make ≠ BinaryHeap in invariant:** OCaml's Set stores unique elements, so re-inserting (alt, v) after a relaxation replaces the old entry automatically. Rust's BinaryHeap allows duplicates, so outdated (old_dist, v) entries pile up and must be filtered with the if d > dist[u] { continue } guard.

**Reverse is the idiomatic min-heap:** Rust's BinaryHeap implements a max-heap (the most common use case for event scheduling). std::cmp::Reverse<T> flips the Ord comparison without any extra crate — it is the canonical, zero-cost way to get a min-heap from BinaryHeap.

Ownership enables in-place mutation: OCaml's purely functional maps return new values on every add; the GC collects the old ones. Rust's HashMap::insert mutates the existing allocation — no allocation per relaxation step — while ownership guarantees there are no aliased references to the old value.

**let else replaces pattern-matching on Option:** Rust's let Some(x) = expr else { return … } is the direct equivalent of OCaml's early-exit if … is_empty then … else let x = …. It keeps the happy path un-indented without needing a full match.

Lifetime annotations encode the OCaml module boundary: OCaml's SMap.Make(String) ties the key type to string structurally. Rust's 'a lifetime ties all &str keys in the graph and the result map to the same borrow scope, preventing dangling references to node names — a guarantee OCaml gets for free from GC.

When to Use Each Style

Use idiomatic Rust when: writing production code, algorithms on large graphs, or any context where readability and debuggability matter. The while let loop is linear in appearance and easy to step through with a debugger.

Use recursive Rust when: demonstrating the direct OCaml translation, teaching the equivalence between tail recursion and loops, or in contexts where the immutable-accumulator mental model aids reasoning about correctness (e.g. property-based testing with fixed state snapshots).

Exercises

Extend the graph to a labeled multigraph where each edge also carries a string label, and return the sequence of edge labels along the shortest path.

Implement Johnson's algorithm: reweight negative edges using Bellman–Ford, then run Dijkstra from every node, enabling all-pairs shortest paths on sparse graphs with negative weights.

Apply Dijkstra to model a packet-routing simulation: nodes are routers, edges are links with latency, and each router table is computed as the shortest-path tree from that router.

Open Source Repos

functional-rust

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

Rust