leetcode_rust/
kth_largest_element_in_an_array.rs

1#![allow(dead_code)]
2
3pub fn find_kth_largest(mut nums: Vec<i32>, k: i32) -> i32 {
4    let k = nums.len() as i32 - k + 1;
5    let mut i = 0;
6    let mut j = nums.len() as i32 - 1;
7    let mut pivot;
8    while i < j {
9        pivot = partition(&mut nums, i, j);
10        if pivot > k - 1 {
11            j = pivot - 1;
12        } else if pivot < k - 1 {
13            i = pivot + 1;
14        } else {
15            break;
16        }
17    }
18    nums[k as usize - 1]
19}
20
21fn partition(nums: &mut Vec<i32>, left: i32, right: i32) -> i32 {
22    let pivot_val = nums[left as usize];
23    let mut i = left as usize;
24    let mut j = right as usize;
25    while i < j {
26        while i < j && nums[j] >= pivot_val {
27            j -= 1;
28        }
29        nums[i] = nums[j];
30
31        while i < j && nums[i] <= pivot_val {
32            i += 1;
33        }
34        nums[j] = nums[i];
35    }
36
37    nums[i] = pivot_val;
38    i as i32
39}
40
41fn partition2(nums: &mut Vec<i32>, left: i32, right: i32) -> i32 {
42    let left = left as usize;
43    let right = right as usize;
44    let pivot_val = nums[left];
45    nums.swap(left as usize, right as usize);
46    let mut store_index = left;
47    for i in left..right {
48        if nums[i] < pivot_val {
49            nums.swap(store_index, i);
50            store_index += 1;
51        }
52    }
53
54    nums.swap(right, store_index);
55    store_index as i32
56}
57
58#[cfg(test)]
59mod tests {
60    use super::*;
61
62    #[test]
63    fn test1() {
64        let mut nums = vec![3, 2, 1, 5, 6, 4];
65        let len = nums.len() as i32 - 1;
66        partition(&mut nums, 0, len);
67        assert_eq!(nums, vec![1, 2, 3, 5, 6, 4]);
68
69        let nums = vec![3, 2, 1, 5, 6, 4];
70        let val = find_kth_largest(nums, 2);
71        assert_eq!(val, 5);
72
73        let nums = vec![-1, -1];
74        let val = find_kth_largest(nums, 2);
75        assert_eq!(val, -1);
76    }
77
78    #[test]
79    fn test2() {
80        let mut nums = vec![3, 2, 1, 4, 5, 6];
81        let len = nums.len() as i32 - 1;
82        partition2(&mut nums, 0, len);
83        assert_eq!(nums, vec![2, 1, 3, 4, 5, 6]);
84    }
85}
leetcode_rust/kth_largest_element_in_an_array.rs

leetcode_rust/
kth_largest_element_in_an_array.rs
