快速选择

快速选择

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

• 时间复杂度 O(n)，空间复杂度 O(1)

#include <vector>
#include <cstdlib>
#include <ctime>

using namespace std;

class Solution {
public:
    int quickSelect(vector<int> &nums, int left, int right, int target) {
        if (left > right) return -1;
        if (left == right) return nums[left];
        int i = left;
        int j = right;
        // 随机化
        srand(time(nullptr));
        int pick = left + (rand() % (right - left));
        swap(nums[left], nums[pick]);
        int pivot = nums[left];

        while (i < j) {
            while (i < j && pivot <= nums[j]) j--;
            if (i < j) nums[i++] = nums[j];
            while (i < j && pivot >= nums[i]) i++;
            if (i < j) nums[j--] = nums[i];
        }
        nums[i] = pivot;
        if (i > target) return quickSelect(nums, left, i - 1, target);
        if (i < target) return quickSelect(nums, i + 1, right, target);
        // i == target
        return nums[i];
    }

    int findKthLargest(vector<int> &nums, int k) {
        // 找第 k 大的就是找递增序列中下标为 nums.size() - k 的
        return quickSelect(nums, 0, nums.size() - 1, nums.size() - k);
    }
};

#include <vector>
#include <cstdlib>
#include <ctime>

using namespace std;

class Solution {
public:
    // 荷兰国旗版
    int quickSelect(vector<int> &nums, int left, int right, int target) {
        if (left > right) return -1;
        if (left == right) return nums[left];
        // 随机化
        srand(time(nullptr));
        int pick = left + (rand() % (right - left));
        swap(nums[left], nums[pick]);
        int pivot = nums[left];

        int l = left;
        int r = right;
        int index = l;
        while (index <= r) {
            if (nums[index] == pivot) {
                index++;
            } else if (nums[index] > pivot) {
                swap(nums[index], nums[r--]);
            } else {
                swap(nums[index++], nums[l++]);
            }
        }

        if (target < l) return quickSelect(nums, left, l - 1, target);
        if (target > r) return quickSelect(nums, r + 1, right, target);
        // 在 pivot 的范围内
        return pivot;
    }

    int findKthLargest(vector<int> &nums, int k) {
        // 找第 k 大的就是找递增序列中下标为 nums.size() - k 的
        return quickSelect(nums, 0, nums.size() - 1, nums.size() - k);
    }
};