use std::cmp::Reverse;
use std::collections::BinaryHeap;

/// Computes shortest distances from `start` using Dijkstra's algorithm.
///
/// The graph is an adjacency list where `graph[u]` contains `(v, weight)` edges.
/// All weights must be non-negative, represented here as `u64`.
pub fn dijkstra(graph: &[Vec<(usize, u64)>], start: usize) -> Vec<Option<u64>> {
    assert!(start < graph.len(), "start node is outside the graph");

    let mut distances = vec![None; graph.len()];
    let mut heap = BinaryHeap::new();

    distances[start] = Some(0);
    heap.push((Reverse(0_u64), start));

    while let Some((Reverse(cost), node)) = heap.pop() {
        // Skip stale heap entries produced before a shorter path was found.
        if distances[node] != Some(cost) {
            continue;
        }

        for &(next, weight) in &graph[node] {
            assert!(next < graph.len(), "edge points outside the graph");

            let Some(next_cost) = cost.checked_add(weight) else {
                continue;
            };

            if distances[next].is_none_or(|current| next_cost < current) {
                distances[next] = Some(next_cost);
                heap.push((Reverse(next_cost), next));
            }
        }
    }

    distances
}

fn main() {
    let graph = vec![
        vec![(1, 4), (2, 1)],
        vec![(3, 1)],
        vec![(1, 2), (3, 5)],
        vec![],
    ];

    println!("{:?}", dijkstra(&graph, 0));
}

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

    #[test]
    fn finds_shortest_paths() {
        let graph = vec![
            vec![(1, 4), (2, 1)],
            vec![(3, 1)],
            vec![(1, 2), (3, 5)],
            vec![],
        ];

        assert_eq!(dijkstra(&graph, 0), vec![Some(0), Some(3), Some(1), Some(4)]);
    }

    #[test]
    fn leaves_unreachable_nodes_empty() {
        let graph = vec![
            vec![(1, 2)],
            vec![],
            vec![],
        ];

        assert_eq!(dijkstra(&graph, 0), vec![Some(0), Some(2), None]);
    }
}