rustsim_pathfinding/

yen.rs

//! Yen's K-shortest loopless paths algorithm.
//!
//! Finds up to K distinct loopless paths between a source and destination node,
//! ordered by ascending cost. Uses Dijkstra as the inner shortest-path routine.
//!
//! The implementation follows Yen (1971) with standard spur-node iteration.
//!
//! # Example
//!
//! ```
//! use rustsim_pathfinding::yen::{yen_k_shortest, YenPath};
//!
//! // Diamond graph: 0→1 (cost 1), 0→2 (cost 2), 1→3 (cost 2), 2→3 (cost 1)
//! let neighbors = |node: &usize| -> Vec<(usize, f64)> {
//!     match *node {
//!         0 => vec![(1, 1.0), (2, 2.0)],
//!         1 => vec![(3, 2.0)],
//!         2 => vec![(3, 1.0)],
//!         _ => vec![],
//!     }
//! };
//!
//! let paths = yen_k_shortest(0, 3, 3, neighbors);
//! assert_eq!(paths.len(), 2);
//! assert_eq!(paths[0].nodes, vec![0, 1, 3]);
//! assert_eq!(paths[1].nodes, vec![0, 2, 3]);
//! ```
//!
//! # References
//!
//! Yen, J. Y. (1971). "Finding the K Shortest Loopless Paths in a Network."
//! Management Science, 17(11), 712–716.

use std::cmp::Ordering;
use std::collections::{BinaryHeap, HashMap, HashSet};

/// A single path with its node sequence and total cost.
#[derive(Debug, Clone)]
pub struct YenPath<N> {
    /// Ordered node sequence from origin to destination (inclusive).
    pub nodes: Vec<N>,
    /// Total path cost.
    pub cost: f64,
}

#[derive(Clone)]
struct DijkEntry<N> {
    node: N,
    cost: f64,
}

impl<N: PartialEq> PartialEq for DijkEntry<N> {
    fn eq(&self, other: &Self) -> bool {
        self.node == other.node
    }
}

impl<N: Eq> Eq for DijkEntry<N> {}

impl<N: Eq> PartialOrd for DijkEntry<N> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<N: Eq> Ord for DijkEntry<N> {
    fn cmp(&self, other: &Self) -> Ordering {
        other
            .cost
            .partial_cmp(&self.cost)
            .unwrap_or(Ordering::Equal)
    }
}

/// Dijkstra shortest path from `src` to `dest`, respecting excluded edges and
/// nodes.
fn dijkstra_path<N, FN, I>(
    src: N,
    dest: N,
    neighbors: &mut FN,
    excluded_edges: &HashSet<(N, N)>,
    excluded_nodes: &HashSet<N>,
) -> Option<YenPath<N>>
where
    N: Clone + Eq + std::hash::Hash,
    FN: FnMut(&N) -> I,
    I: IntoIterator<Item = (N, f64)>,
{
    let mut g_scores: HashMap<N, f64> = HashMap::new();
    let mut came_from: HashMap<N, N> = HashMap::new();
    let mut closed: HashSet<N> = HashSet::new();
    let mut heap: BinaryHeap<DijkEntry<N>> = BinaryHeap::new();

    g_scores.insert(src.clone(), 0.0);
    heap.push(DijkEntry {
        node: src.clone(),
        cost: 0.0,
    });

    while let Some(current) = heap.pop() {
        if current.node == dest {
            // Reconstruct path.
            let mut path = Vec::new();
            let mut cur = dest.clone();
            loop {
                path.push(cur.clone());
                match came_from.get(&cur) {
                    Some(prev) => cur = prev.clone(),
                    None => break,
                }
            }
            path.reverse();
            return Some(YenPath {
                nodes: path,
                cost: current.cost,
            });
        }

        if !closed.insert(current.node.clone()) {
            continue;
        }

        let g = current.cost;
        for (nbr, edge_cost) in neighbors(&current.node) {
            if closed.contains(&nbr) {
                continue;
            }
            if excluded_nodes.contains(&nbr) {
                continue;
            }
            if excluded_edges.contains(&(current.node.clone(), nbr.clone())) {
                continue;
            }
            let tentative = g + edge_cost;
            let prev = g_scores.get(&nbr).copied().unwrap_or(f64::INFINITY);
            if tentative < prev {
                g_scores.insert(nbr.clone(), tentative);
                came_from.insert(nbr.clone(), current.node.clone());
                heap.push(DijkEntry {
                    node: nbr,
                    cost: tentative,
                });
            }
        }
    }

    None
}

/// Find up to `k` shortest loopless paths from `src` to `dest`.
///
/// Uses Yen's algorithm with Dijkstra as the inner shortest-path routine.
///
/// # Arguments
///
/// - `src` - source node
/// - `dest` - destination node
/// - `k` - maximum number of paths to find
/// - `neighbors` - closure returning (neighbor, edge_cost) pairs for a node
///
/// Returns paths sorted by ascending cost. May return fewer than `k` paths
/// if fewer distinct loopless paths exist.
pub fn yen_k_shortest<N, FN, I>(src: N, dest: N, k: usize, mut neighbors: FN) -> Vec<YenPath<N>>
where
    N: Clone + Eq + std::hash::Hash,
    FN: FnMut(&N) -> I,
    I: IntoIterator<Item = (N, f64)>,
{
    if k == 0 {
        return Vec::new();
    }

    // Find the shortest path first.
    let first = dijkstra_path(
        src.clone(),
        dest.clone(),
        &mut neighbors,
        &HashSet::new(),
        &HashSet::new(),
    );
    let Some(first) = first else {
        return Vec::new();
    };

    let mut accepted: Vec<YenPath<N>> = vec![first];

    // Candidate heap: (cost, path_nodes).
    // We wrap in a struct for ordering.
    let mut candidates: BinaryHeap<CandidateEntry<N>> = BinaryHeap::new();
    let mut candidate_set: HashSet<Vec<N>> = HashSet::new();

    for ki in 1..k {
        let prev_path = &accepted[ki - 1].nodes;

        // Spur from each node along the previous path (except dest).
        for spur_idx in 0..prev_path.len().saturating_sub(1) {
            let spur_node = prev_path[spur_idx].clone();
            let root_path: Vec<N> = prev_path[..=spur_idx].to_vec();

            // Exclude edges at spur_node that share the same root prefix.
            let mut excluded_edges: HashSet<(N, N)> = HashSet::new();
            for accepted_path in &accepted {
                if accepted_path.nodes.len() > spur_idx
                    && accepted_path.nodes[..=spur_idx] == root_path[..]
                {
                    excluded_edges.insert((
                        accepted_path.nodes[spur_idx].clone(),
                        accepted_path.nodes[spur_idx + 1].clone(),
                    ));
                }
            }

            // Exclude nodes in root (except spur node) to guarantee loopless.
            let mut excluded_nodes: HashSet<N> = HashSet::new();
            for node in &root_path[..spur_idx] {
                excluded_nodes.insert(node.clone());
            }

            if let Some(spur_path) = dijkstra_path(
                spur_node,
                dest.clone(),
                &mut neighbors,
                &excluded_edges,
                &excluded_nodes,
            ) {
                // Concatenate root + spur (spur starts at spur_node).
                let mut full_nodes = root_path.clone();
                full_nodes.extend_from_slice(&spur_path.nodes[1..]);

                // Check for loops.
                let mut seen = HashSet::new();
                if full_nodes.iter().any(|n| !seen.insert(n.clone())) {
                    continue;
                }

                if !candidate_set.contains(&full_nodes) {
                    // Compute total cost.
                    let cost = path_cost(&full_nodes, &mut neighbors);
                    candidate_set.insert(full_nodes.clone());
                    candidates.push(CandidateEntry {
                        cost,
                        nodes: full_nodes,
                    });
                }
            }
        }

        // Pop the cheapest candidate.
        if let Some(best) = candidates.pop() {
            accepted.push(YenPath {
                nodes: best.nodes,
                cost: best.cost,
            });
        } else {
            break;
        }
    }

    accepted
}

/// Compute the total cost of a path by summing edge costs from the neighbor
/// function.
fn path_cost<N, FN, I>(nodes: &[N], neighbors: &mut FN) -> f64
where
    N: Clone + Eq + std::hash::Hash,
    FN: FnMut(&N) -> I,
    I: IntoIterator<Item = (N, f64)>,
{
    let mut total = 0.0;
    for pair in nodes.windows(2) {
        let from = &pair[0];
        let to = &pair[1];
        // Find the edge cost from `from` to `to`.
        let edge_cost = neighbors(from)
            .into_iter()
            .find(|(n, _)| n == to)
            .map(|(_, c)| c)
            .unwrap_or(0.0);
        total += edge_cost;
    }
    total
}

struct CandidateEntry<N> {
    cost: f64,
    nodes: Vec<N>,
}

impl<N: PartialEq> PartialEq for CandidateEntry<N> {
    fn eq(&self, other: &Self) -> bool {
        self.cost == other.cost && self.nodes == other.nodes
    }
}

impl<N: Eq> Eq for CandidateEntry<N> {}

impl<N: Eq> PartialOrd for CandidateEntry<N> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<N: Eq> Ord for CandidateEntry<N> {
    fn cmp(&self, other: &Self) -> Ordering {
        // Min-heap: reverse the cost comparison.
        other
            .cost
            .partial_cmp(&self.cost)
            .unwrap_or(Ordering::Equal)
    }
}

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

    fn diamond_neighbors(node: &usize) -> Vec<(usize, f64)> {
        match *node {
            0 => vec![(1, 1.0), (2, 2.0)],
            1 => vec![(3, 2.0)],
            2 => vec![(3, 1.0)],
            _ => vec![],
        }
    }

    #[test]
    fn finds_two_paths_on_diamond() {
        let paths = yen_k_shortest(0, 3, 3, diamond_neighbors);
        assert_eq!(paths.len(), 2);
        assert_eq!(paths[0].nodes, vec![0, 1, 3]);
        assert!((paths[0].cost - 3.0).abs() < 1e-6);
        assert_eq!(paths[1].nodes, vec![0, 2, 3]);
        assert!((paths[1].cost - 3.0).abs() < 1e-6);
    }

    #[test]
    fn single_path_on_line() {
        let neighbors = |node: &usize| -> Vec<(usize, f64)> {
            match *node {
                0 => vec![(1, 1.0)],
                1 => vec![(2, 1.0)],
                _ => vec![],
            }
        };
        let paths = yen_k_shortest(0, 2, 5, neighbors);
        assert_eq!(paths.len(), 1);
        assert_eq!(paths[0].nodes, vec![0, 1, 2]);
        assert!((paths[0].cost - 2.0).abs() < 1e-6);
    }

    #[test]
    fn no_path_returns_empty() {
        let neighbors = |_: &usize| -> Vec<(usize, f64)> { vec![] };
        let paths = yen_k_shortest(0, 5, 3, neighbors);
        assert!(paths.is_empty());
    }

    #[test]
    fn k_zero_returns_empty() {
        let paths = yen_k_shortest(0, 3, 0, diamond_neighbors);
        assert!(paths.is_empty());
    }

    #[test]
    fn paths_are_loopless() {
        // Graph with potential loops: 0→1→2→3, 0→2→3, 0→1→0 (loop edge)
        let neighbors = |node: &usize| -> Vec<(usize, f64)> {
            match *node {
                0 => vec![(1, 1.0), (2, 3.0)],
                1 => vec![(0, 1.0), (2, 1.0)],
                2 => vec![(3, 1.0)],
                _ => vec![],
            }
        };
        let paths = yen_k_shortest(0, 3, 5, neighbors);
        for path in &paths {
            let mut seen = HashSet::new();
            assert!(
                path.nodes.iter().all(|n| seen.insert(n)),
                "Path contains loop: {:?}",
                path.nodes
            );
        }
    }

    #[test]
    fn paths_sorted_by_cost() {
        // Grid-like graph with multiple paths
        let neighbors = |node: &usize| -> Vec<(usize, f64)> {
            match *node {
                0 => vec![(1, 1.0), (2, 2.0), (3, 5.0)],
                1 => vec![(4, 1.0)],
                2 => vec![(4, 1.0)],
                3 => vec![(4, 1.0)],
                _ => vec![],
            }
        };
        let paths = yen_k_shortest(0, 4, 5, neighbors);
        for i in 1..paths.len() {
            assert!(
                paths[i].cost >= paths[i - 1].cost - 1e-12,
                "Paths not sorted: cost[{}]={} < cost[{}]={}",
                i,
                paths[i].cost,
                i - 1,
                paths[i - 1].cost
            );
        }
    }
}