diff --git a/rs/cheapest-flights-within-k-stops.rs b/rs/cheapest-flights-within-k-stops.rs new file mode 100644 index 0000000..d4733e8 --- /dev/null +++ b/rs/cheapest-flights-within-k-stops.rs @@ -0,0 +1,108 @@ +use std::cmp; +use std::collections::BinaryHeap; +use std::convert::TryInto; + +struct MinHeap { + heap: BinaryHeap>, +} + +impl MinHeap { + 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 { + 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>, 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 = 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›") + } +}