LeetCode/rs/k-th-smallest-prime-fraction.rs

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

#[derive(Eq)]
struct Fraction {
    num: i32,
    den: i32,
    num_idx: usize,
    den_idx: usize,
}

impl Fraction {
    fn new(primes: &[i32], num_idx: usize, den_idx: usize) -> Self {
        Self {
            num: primes[num_idx],
            den: primes[den_idx],
            num_idx,
            den_idx,
        }
    }
}

impl Ord for Fraction {
    fn cmp(&self, rhs: &Fraction) -> cmp::Ordering {
        let left = self.num * rhs.den;
        let right = rhs.num * self.den;

        left.cmp(&right)
    }
}

impl PartialOrd for Fraction {
    fn partial_cmp(&self, rhs: &Fraction) -> Option<cmp::Ordering> {
        Some(self.cmp(rhs))
    }
}

impl PartialEq for Fraction {
    fn eq(&self, rhs: &Fraction) -> bool {
        self.num * rhs.den == rhs.num * self.den
    }
}

impl From<Fraction> for Vec<i32> {
    fn from(f: Fraction) -> Self {
        vec![f.num, f.den]
    }
}

struct Solution {}
impl Solution {
    pub fn kth_smallest_prime_fraction(arr: Vec<i32>, k: i32) -> Vec<i32> {
        let mut heap: BinaryHeap<cmp::Reverse<Fraction>> = BinaryHeap::new();

        let last = arr.len() - 1;
        for i in 0..arr.len() {
            heap.push(cmp::Reverse(Fraction::new(&arr, i, last)));
        }

        for _ in 1..k {
            let smallest = heap.pop().expect("There is always at least one fraction").0;
            if smallest.den_idx - 1 > smallest.num_idx {
                heap.push(cmp::Reverse(Fraction::new(
                    &arr,
                    smallest.num_idx,
                    smallest.den_idx - 1,
                )));
            }
        }

        heap.pop().expect("There is always k-th smallest").0.into()
    }
}

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

    #[test]
    fn example_1() {
        assert_eq!(
            Solution::kth_smallest_prime_fraction(vec![1, 2, 3, 5], 3),
            vec![2, 5]
        );
    }

    #[test]
    fn example_2() {
        assert_eq!(
            Solution::kth_smallest_prime_fraction(vec![1, 7], 1),
            vec![1, 7]
        );
    }
}