LeetCode 215. 数组中的第K个最大元素

题目链接：https://leetcode.cn/problems/kth-largest-element-in-an-array/

在 快速选择算法 中，我们约定：数组下标为 $[0,n-1]$，对于从小到大排好序的数组来说，第 $0$ 个数就是第 $0$ 小的数，第 $1$ 个数就是第 $1$ 小的数，……，第 $k$ 个数就是第 $k$ 小的数。也就是说，这里的 $k$ 正好和数组下标 $k$ 相对应。

在本题中，数组下标为 $[0,n-1]$，对于从小到大排好序的数组来说，第 $n-1$ 个数就是第 $1$ 大的数，第 $n-2$ 个数就是第 $2$ 大的数，……，第 $n-k$ 个数就是第 $k$ 大的数。也就是说，本题的 $k$ 正好和数组下标 $n-k$ 相对应。

class Solution {
public:
    // 在[left,right]中寻找第k小的数
    int quickSelect(vector<int>& nums, int left, int right, int k) {
        if (left >= right) {
            return nums[left];
        }

        // 取最左边的元素为轴点元素
        int i = left, j = right, pivot = nums[left];

        while (i < j) {
            while (i < j && nums[j] > pivot) {
                j--;
            }
            if (i < j) {
                nums[i++] = nums[j];
            }

            while (i < j && nums[i] < pivot) {
                i++;
            }
            if (i < j) {
                nums[j--] = nums[i];
            }
        }

        // while循环退出时，i就是轴点位置，将轴点元素归位
        nums[i] = pivot;

        if (k == i) {
            // 如果轴点元素就是第k小的数，则直接返回轴点元素
            return nums[i];
        } else if (k < i) {
            // 在轴点元素的左边[left,i-1]中寻找第k小的数
            return quickSelect(nums, left, i - 1, k);
        } else {
            // 在轴点元素的右边[i+1,right]中寻找第k小的数
            return quickSelect(nums, i + 1, right, k);
        }
    }

    int findKthLargest(vector<int>& nums, int k) {
        int n = nums.size();
        return quickSelect(nums, 0, n - 1, n - k);
    }
};