查找第K小的元素

  1. use crate::sorting::partition;
  2. use std::cmp::{Ordering, PartialOrd};
  4. /// Returns k-th smallest element of an array, i.e. its order statistics.
  5. /// Time complexity is O(n^2) in the worst case, but only O(n) on average.
  6. /// It mutates the input, and therefore does not require additional space.
  7. pub fn kth_smallest<T>(input: &mut [T], k: usize) -> Option<T>
  8. where
  9.     T: PartialOrd + Copy,
  10. {
  11.     if input.is_empty() {
  12.         return None;
  13.     }
  15.     let kth = _kth_smallest(input, k, 0, input.len() - 1);
  16.     Some(kth)
  17. }
  19. fn _kth_smallest<T>(input: &mut [T], k: usize, lo: usize, hi: usize) -> T
  20. where
  21.     T: PartialOrd + Copy,
  22. {
  23.     if lo == hi {
  24.         return input[lo];
  25.     }
  27.     let pivot = partition(input, lo as isize, hi as isize) as usize;
  28.     let i = pivot - lo + 1;
  30.     match k.cmp(&i) {
  31.         Ordering::Equal => input[pivot],
  32.         Ordering::Less => _kth_smallest(input, k, lo, pivot - 1),
  33.         Ordering::Greater => _kth_smallest(input, k - i, pivot + 1, hi),
  34.     }
  35. }
  37. #[cfg(test)]
  38. mod tests {
  39.     use super::*;
  41.     #[test]
  42.     fn empty() {
  43.         let mut zero: [u8; 0] = [];
  44.         let first = kth_smallest(&mut zero, 1);
  46.         assert_eq!(None, first);
  47.     }
  49.     #[test]
  50.     fn one_element() {
  51.         let mut one = [1];
  52.         let first = kth_smallest(&mut one, 1);
  54.         assert_eq!(1, first.unwrap());
  55.     }
  57.     #[test]
  58.     fn many_elements() {
  59.         // 0 1 3 4 5 7 8 9 9 10 12 13 16 17
  60.         let mut many = [9, 17, 3, 16, 13, 10, 1, 5, 7, 12, 4, 8, 9, 0];
  62.         let first = kth_smallest(&mut many, 1);
  63.         let third = kth_smallest(&mut many, 3);
  64.         let sixth = kth_smallest(&mut many, 6);
  65.         let fourteenth = kth_smallest(&mut many, 14);
  67.         assert_eq!(0, first.unwrap());
  68.         assert_eq!(3, third.unwrap());
  69.         assert_eq!(7, sixth.unwrap());
  70.         assert_eq!(17, fourteenth.unwrap());
  71.     }
  72. }