108 lines
2.5 KiB
Rust
108 lines
2.5 KiB
Rust
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›")
|
||
}
|
||
}
|