LeetCode/rs/cheapest-flights-within-k-stops.rs

use std::cmp;
use std::collections::BinaryHeap;
use std::convert::TryInto;

struct MinHeap<T: Ord> {
    heap: BinaryHeap<cmp::Reverse<T>>,
}

impl<T: Ord> MinHeap<T> {
    pub fn new() -> Self {
        Self {
            heap: BinaryHeap::new(),
        }
    }

    pub fn push(&mut self, item: T) {
        self.heap.push(cmp::Reverse(item));
    }

    pub fn pop(&mut self) -> Option<T> {
        self.heap.pop().map(|cmp::Reverse(item)| item)
    }

    pub fn is_empty(&self) -> bool {
        self.heap.is_empty()
    }

    pub fn peek(&self) -> Option<&T> {
        self.heap.peek().map(|cmp::Reverse(item)| item)
    }
}

#[derive(Debug, Clone, Copy)]
struct DijkstraEntry {
    price: i64,
    stops: i32,
}

impl DijkstraEntry {
    fn new() -> Self {
        Self {
            price: i64::MAX,
            stops: i32::MAX,
        }
    }
}

impl Solution {
    pub fn find_cheapest_price(n: i32, flights: Vec<Vec<i32>>, src: i32, dst: i32, k: i32) -> i32 {
        let n = n as usize;
        let src = src as usize;
        let dst = dst as usize;

        let graph = {
            let mut g: Vec<i64> = vec![i32::MAX.into(); n * n];

            for flight in &flights {
                let u = flight[0] as usize;
                let v = flight[1] as usize;
                let w = flight[2].into();

                g[u * n + v] = w;
            }

            g
        };

        let mut data = vec![DijkstraEntry::new(); n];
        let mut q: MinHeap<(i32, i64, usize)> = MinHeap::new();

        // initialize for the first airport
        q.push((0, 0, src));

        // run Dijkstra
        while let Some((stops, price, idx)) = q.pop() {
            println!("Entry ({price}, {stops}, {idx})");

            // higher price
            if price > data[idx].price {
                println!("  Skipping because of price");
                continue;
            }

            data[idx].price = price;
            data[idx].stops = stops;

            // can't fly further
            if stops > k {
                println!("  Can't fly further");
                continue;
            }

            for next in 0..n {
                let flight_price = graph[idx * n + next];
                if flight_price == i32::MAX as i64 {
                    continue;
                }

                q.push((stops + 1, price + flight_price, next));
            }
        }

        if data[dst].price == i64::MAX {
            return -1;
        }
        data[dst].price.try_into().expect("price should fit ‹i32›")
    }
}