ExamplesBy LevelBy TopicLearning Paths
1147 Advanced
← 11461148 →

Dijkstra's Shortest Path — Functional Priority Queue

GraphsAlgorithms

Tutorial

The Problem

Find the shortest path from a start node to all reachable nodes in a weighted directed graph. This example emphasises the functional priority-queue idiom: OCaml uses Set.Make as an ordered-set priority queue; Rust mirrors this with BTreeSet.

🎯 Learning Outcomes

  • • How OCaml's Set.Make doubles as a min-heap via ordered-set semantics
  • • How Rust's BTreeSet<(usize, String)> is the structural equivalent
  • • How BTreeMap mirrors OCaml's Map.Make(String) — deterministic sorted output
  • • How OCaml's let rec go pq dist tail-recursive loop translates to a Rust fn go(…) recursive helper with fold over neighbours

    • 🦀 The Rust Way

    BTreeSet<(usize, String)> is Rust's structural analogue: tuples compare lexicographically, so the minimum-distance entry is always first. iter().next() is PQ.min_elt; remove is PQ.remove. BTreeMap<String, usize> mirrors Map.Make(String) and produces sorted output without an explicit sort step. A second implementation uses a recursive go helper with Iterator::fold over neighbours, matching the OCaml List.fold_left pattern directly.

    Code Example

    use std::collections::{BTreeMap, BTreeSet};

type Graph = BTreeMap<String, Vec<(String, usize)>>;

fn dijkstra(graph: &Graph, start: &str) -> BTreeMap<String, usize> {
    let mut dist: BTreeMap<String, usize> = BTreeMap::from([(start.to_string(), 0)]);
    let mut pq: BTreeSet<(usize, String)> = BTreeSet::from([(0, start.to_string())]);

    while let Some((d, u)) = pq.iter().next().cloned() {
        pq.remove(&(d, u.clone()));
        if dist.get(&u).is_some_and(|&best| d > best) { continue; }
        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(&usize::MAX) {
                dist.insert(v.clone(), alt);
                pq.insert((alt, v.clone()));
            }
        }
    }
    dist
}

    Key Differences

  • Priority queue: OCaml uses Set.Make (functional BST); Rust uses BTreeSet (mutable B-tree). Both give O(log n) min-extraction and O(log n) insertion.
  • Distance map: OCaml's Map.Make(String) is persistent/immutable; Rust's BTreeMap is mutable but produces the same sorted-key iteration order.
  • Tail recursion: OCaml's TCO turns let rec go into a loop automatically; Rust's fn go recurses on the stack — safe for small graphs, not for production-scale graphs.
  • Stale entries: Both approaches allow duplicate (dist, node) entries in the priority queue; the alt < current guard ensures only improvements are enqueued.

    • OCaml Approach

    OCaml uses Set.Make with a custom comparator on (int * string) tuples as a functional priority queue. PQ.min_elt extracts the minimum cheaply; PQ.remove shadows the binding to produce a new set. Map.Make(String) holds distances immutably. The inner let rec go pq dist is a tail-recursive loop that the OCaml compiler optimises into iteration.

    Full Source

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

/// Graph keyed by node name; each node maps to its (neighbour, weight) list.
/// BTreeMap mirrors OCaml's `Map.Make(String)` — sorted, deterministic iteration.
pub type Graph = BTreeMap<String, Vec<(String, usize)>>;

/// Dijkstra using `BTreeSet<(usize, String)>` as a priority queue.
///
/// This mirrors OCaml's `Set.Make` idiom: an ordered set doubles as a min-heap.
/// `BTreeSet::iter().next()` gives the minimum element in O(log n), just like
/// OCaml's `PQ.min_elt`. `BTreeSet::remove` mirrors `PQ.remove`.
///
/// `BTreeMap` for distances gives sorted output like OCaml's `Map.Make(String)`.
pub fn dijkstra(graph: &Graph, start: &str) -> BTreeMap<String, usize> {
    let mut dist: BTreeMap<String, usize> = BTreeMap::from([(start.to_string(), 0)]);
    let mut pq: BTreeSet<(usize, String)> = BTreeSet::from([(0, start.to_string())]);

    while let Some((d, u)) = pq.iter().next().cloned() {
        // Mirrors OCaml: let pq = PQ.remove (d, u) pq
        pq.remove(&(d, u.clone()));

        // Stale entry: a shorter path to u was already finalised
        if dist.get(&u).is_some_and(|&best| d > best) {
            continue;
        }

        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(&usize::MAX) {
                dist.insert(v.clone(), alt);
                // Mirrors OCaml: PQ.add (alt, v) pq
                pq.insert((alt, v.clone()));
            }
        }
    }

    dist
}

/// Dijkstra — recursive functional style.
///
/// Mirrors OCaml's `let rec go pq dist = if PQ.is_empty pq then dist else ...`
/// and `List.fold_left` over neighbours.
///
/// # Note
/// Rust does not guarantee tail-call optimisation; deep graphs may stack-overflow.
/// Use `dijkstra` for production use; this variant demonstrates the OCaml parallel.
pub fn dijkstra_recursive(graph: &Graph, start: &str) -> BTreeMap<String, usize> {
    let dist = BTreeMap::from([(start.to_string(), 0)]);
    let pq = BTreeSet::from([(0usize, start.to_string())]);
    go(graph, pq, dist)
}

/// Recursive worker — mirrors OCaml's `let rec go pq dist = ...`
fn go(
    graph: &Graph,
    mut pq: BTreeSet<(usize, String)>,
    dist: BTreeMap<String, usize>,
) -> BTreeMap<String, usize> {
    // if PQ.is_empty pq then dist
    let Some((d, u)) = pq.iter().next().cloned() else {
        return dist;
    };
    // let pq = PQ.remove (d, u) pq
    pq.remove(&(d, u.clone()));

    let neighbors = graph.get(&u).map(Vec::as_slice).unwrap_or(&[]);

    // Mirrors: List.fold_left (fun (dist, pq) (v, w) -> ...) (dist, pq) neighbors
    let (dist, pq) = neighbors
        .iter()
        .fold((dist, pq), |(mut dist, mut pq), (v, w)| {
            let alt = d + w;
            let current = *dist.get(v).unwrap_or(&usize::MAX);
            if alt < current {
                dist.insert(v.clone(), alt);
                pq.insert((alt, v.clone()));
            }
            (dist, pq)
        });

    go(graph, pq, dist)
}

/// Build a `Graph` from a slice of `(from, [(to, weight)])` pairs — convenience for tests.
pub fn build_graph(edges: &[(&str, &[(&str, usize)])]) -> Graph {
    edges
        .iter()
        .map(|(from, neighbors)| {
            let ns = neighbors
                .iter()
                .map(|(to, w)| ((*to).to_string(), *w))
                .collect();
            ((*from).to_string(), ns)
        })
        .collect()
}

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

    fn sample_graph() -> Graph {
        build_graph(&[
            ("a", &[("b", 1), ("c", 4)]),
            ("b", &[("c", 2), ("d", 6)]),
            ("c", &[("d", 3)]),
            ("d", &[]),
        ])
    }

    #[test]
    fn test_dijkstra_distances_from_a() {
        let dist = dijkstra(&sample_graph(), "a");
        assert_eq!(dist["a"], 0);
        assert_eq!(dist["b"], 1);
        assert_eq!(dist["c"], 3); // a->b->c = 1+2=3, shorter than a->c = 4
        assert_eq!(dist["d"], 6); // a->b->c->d = 1+2+3=6
    }

    #[test]
    fn test_recursive_matches_iterative() {
        let g = sample_graph();
        assert_eq!(dijkstra(&g, "a"), dijkstra_recursive(&g, "a"));
    }

    #[test]
    fn test_single_node_graph() {
        let g = build_graph(&[("x", &[])]);
        let dist = dijkstra(&g, "x");
        assert_eq!(dist["x"], 0);
        assert_eq!(dist.len(), 1);
    }

    #[test]
    fn test_unreachable_nodes_absent() {
        let g = build_graph(&[
            ("a", &[("b", 5)]),
            ("c", &[("d", 2)]), // disconnected component
            ("b", &[]),
            ("d", &[]),
        ]);
        let dist = dijkstra(&g, "a");
        assert_eq!(dist.get("a"), Some(&0));
        assert_eq!(dist.get("b"), Some(&5));
        assert_eq!(dist.get("c"), None); // unreachable from "a"
        assert_eq!(dist.get("d"), None);
    }

    #[test]
    fn test_multiple_paths_picks_shortest() {
        // Diamond: a->b (1), a->c (10), b->d (1), c->d (1)
        // Shortest to d: a->b->d = 2, not a->c->d = 11
        let g = build_graph(&[
            ("a", &[("b", 1), ("c", 10)]),
            ("b", &[("d", 1)]),
            ("c", &[("d", 1)]),
            ("d", &[]),
        ]);
        let dist = dijkstra(&g, "a");
        assert_eq!(dist["d"], 2);
    }

    #[test]
    fn test_btreemap_output_is_sorted() {
        // BTreeMap mirrors OCaml's Map.Make — iteration order is always sorted
        let dist = dijkstra(&sample_graph(), "a");
        let keys: Vec<_> = dist.keys().cloned().collect();
        let mut expected = keys.clone();
        expected.sort();
        assert_eq!(keys, expected);
    }
}
    ✓ Tests Rust test suite 
    #[cfg(test)]
mod tests {
    use super::*;

    fn sample_graph() -> Graph {
        build_graph(&[
            ("a", &[("b", 1), ("c", 4)]),
            ("b", &[("c", 2), ("d", 6)]),
            ("c", &[("d", 3)]),
            ("d", &[]),
        ])
    }

    #[test]
    fn test_dijkstra_distances_from_a() {
        let dist = dijkstra(&sample_graph(), "a");
        assert_eq!(dist["a"], 0);
        assert_eq!(dist["b"], 1);
        assert_eq!(dist["c"], 3); // a->b->c = 1+2=3, shorter than a->c = 4
        assert_eq!(dist["d"], 6); // a->b->c->d = 1+2+3=6
    }

    #[test]
    fn test_recursive_matches_iterative() {
        let g = sample_graph();
        assert_eq!(dijkstra(&g, "a"), dijkstra_recursive(&g, "a"));
    }

    #[test]
    fn test_single_node_graph() {
        let g = build_graph(&[("x", &[])]);
        let dist = dijkstra(&g, "x");
        assert_eq!(dist["x"], 0);
        assert_eq!(dist.len(), 1);
    }

    #[test]
    fn test_unreachable_nodes_absent() {
        let g = build_graph(&[
            ("a", &[("b", 5)]),
            ("c", &[("d", 2)]), // disconnected component
            ("b", &[]),
            ("d", &[]),
        ]);
        let dist = dijkstra(&g, "a");
        assert_eq!(dist.get("a"), Some(&0));
        assert_eq!(dist.get("b"), Some(&5));
        assert_eq!(dist.get("c"), None); // unreachable from "a"
        assert_eq!(dist.get("d"), None);
    }

    #[test]
    fn test_multiple_paths_picks_shortest() {
        // Diamond: a->b (1), a->c (10), b->d (1), c->d (1)
        // Shortest to d: a->b->d = 2, not a->c->d = 11
        let g = build_graph(&[
            ("a", &[("b", 1), ("c", 10)]),
            ("b", &[("d", 1)]),
            ("c", &[("d", 1)]),
            ("d", &[]),
        ]);
        let dist = dijkstra(&g, "a");
        assert_eq!(dist["d"], 2);
    }

    #[test]
    fn test_btreemap_output_is_sorted() {
        // BTreeMap mirrors OCaml's Map.Make — iteration order is always sorted
        let dist = dijkstra(&sample_graph(), "a");
        let keys: Vec<_> = dist.keys().cloned().collect();
        let mut expected = keys.clone();
        expected.sort();
        assert_eq!(keys, expected);
    }
}

    Deep Comparison

    OCaml vs Rust: Dijkstra's Shortest Path — 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 — BTreeSet priority queue)

    use std::collections::{BTreeMap, BTreeSet};

type Graph = BTreeMap<String, Vec<(String, usize)>>;

fn dijkstra(graph: &Graph, start: &str) -> BTreeMap<String, usize> {
    let mut dist: BTreeMap<String, usize> = BTreeMap::from([(start.to_string(), 0)]);
    let mut pq: BTreeSet<(usize, String)> = BTreeSet::from([(0, start.to_string())]);

    while let Some((d, u)) = pq.iter().next().cloned() {
        pq.remove(&(d, u.clone()));
        if dist.get(&u).is_some_and(|&best| d > best) { continue; }
        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(&usize::MAX) {
                dist.insert(v.clone(), alt);
                pq.insert((alt, v.clone()));
            }
        }
    }
    dist
}

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

    fn dijkstra_recursive(graph: &Graph, start: &str) -> BTreeMap<String, usize> {
    let dist = BTreeMap::from([(start.to_string(), 0)]);
    let pq = BTreeSet::from([(0usize, start.to_string())]);
    go(graph, pq, dist)
}

fn go(
    graph: &Graph,
    mut pq: BTreeSet<(usize, String)>,
    dist: BTreeMap<String, usize>,
) -> BTreeMap<String, usize> {
    let Some((d, u)) = pq.iter().next().cloned() else { return dist; };
    pq.remove(&(d, u.clone()));
    let neighbors = graph.get(&u).map(Vec::as_slice).unwrap_or(&[]);
    // Mirrors: List.fold_left (fun (dist, pq) (v, w) -> ...) (dist, pq) neighbors
    let (dist, pq) = neighbors
        .iter()
        .fold((dist, pq), |(mut dist, mut pq), (v, w)| {
            let alt = d + w;
            let current = *dist.get(v).unwrap_or(&usize::MAX);
            if alt < current {
                dist.insert(v.clone(), alt);
                pq.insert((alt, v.clone()));
            }
            (dist, pq)
        });
    go(graph, pq, dist)
}

    Type Signatures

    ConceptOCamlRust
    Priority queue typePQ.t (= Set.Make(int * string))BTreeSet<(usize, String)>
    Distance map typeint SMap.t (= Map.Make(String))BTreeMap<String, usize>
    Graph type(string * int) list SMap.tBTreeMap<String, Vec<(String, usize)>>
    Dijkstra signatureval dijkstra : graph -> string -> int SMap.tfn dijkstra(graph: &Graph, start: &str) -> BTreeMap<String, usize>
    Recursive workerlet rec go : PQ.t -> int SMap.t -> int SMap.tfn go(graph: &Graph, pq: BTreeSet<…>, dist: BTreeMap<…>) -> BTreeMap<…>

    Key Insights

  • **Set.Make ≈ BTreeSet**: OCaml's ordered-set functor gives min-extraction

    • in O(log n) via min_elt; BTreeSet::iter().next() achieves the same because B-trees store entries in sorted order. Both structures treat (dist, node) tuples as the unique key, so the same node can appear at multiple distances — stale entries are naturally superseded.

  • Persistent vs mutable maps: OCaml's Map.Make is purely functional — each

    • SMap.add returns a new map without touching the old one. Rust's BTreeMap is mutable. However, in the recursive go the maps are moved by value, so the ownership semantics produce the same single-threaded semantics: at any point only one version of dist/pq is live.

  • **List.fold_left ≈ Iterator::fold**: The OCaml fold_left (fun (d,p) (v,w) -> ...) (dist,pq) neighbors

    • translates almost character-for-character into Rust's .fold((dist, pq), |(mut dist, mut pq), (v, w)| {...}). Both accumulate (dist, pq) as a pair and return the updated pair.

  • Tail-call optimisation: OCaml's compiler turns let rec go … = … go pq dist

    • into a loop (TCO guaranteed in many cases). Rust makes no such guarantee — fn go recurses on the stack. For small inputs this is fine; for large graphs, convert to a while loop (the dijkstra function above).

  • Sorted output for free: Both Map.Make(String) and BTreeMap iterate in

    • ascending key order — no post-processing needed to get deterministic output, unlike HashMap which would require an explicit .sort().

    When to Use Each Style

    Use the iterative BTreeSet style when: you want a faithful structural mirror of the OCaml Set.Make idiom and sorted output, but need a safe iterative loop that handles graphs of arbitrary depth without stack risk.

    Use the recursive fold style when: demonstrating the direct OCaml→Rust translation for teaching, or when graphs are small and the List.fold_left parallel is the primary learning goal.

    Exercises

  • Generalize the graph representation to use a generic edge-weight type W: Ord + Add + Zero so the same Dijkstra function works on both integer and floating-point graphs.
  • Implement eccentricity — the maximum shortest-path distance from a given node to any reachable node — and use it to compute the graph's diameter and radius.
  • Modify Dijkstra to return a Stream-like iterator that yields (node, distance) pairs in discovery order, enabling early termination once a target node is reached.

    • Open Source Repos

    functional-rust

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

    Rust