查找第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. }