215-数组中的第k个最大元素

堆

LeetCode

发布日期: 2021-05-29

LeetCode 215. Kth Largest Element in an Array

Given an integer array nums and an integer k, return the kth largest element in the array.

Note that it is the kth largest element in the sorted order, not the kth distinct element.

Example 1:

Input: nums = [3,2,1,5,6,4], k = 2
Output: 5

Example 2:

Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
Output: 4

method 1: 小根堆

维护一个小根堆（从小到大排序），如果新的数比堆顶元素大，就把堆顶元素弹出，新的数入队
如果队列是|4|6|，4是第二大的数，后面又来个5，那4就不是第二大的数了，所以把4弹出，5入队，至于6和5谁是第二大的数，优先队列按从小到大排序之后就知道了

如果是问第k个最小的元素，就用大根堆

注意：
小根堆的写法
priority_queue<int, vector<int>, greater<int>>
同理，大根堆
priority_queue<int, vector<int>, less<int>>

int findKthLargest(vector<int>& nums, int k) {
    priority_queue<int, vector<int>, greater<int>> q;
    for (int i = 0; i < nums.size(); ++i) {
        if (i < k) q.push(nums[i]); // 小于k就直接放进来
        else if (nums[i] > q.top()) {
            q.pop();
            q.push(nums[i]);
        }
    }
    return q.top();
}

也可以先把元素放进来，再看有没有大于k，再弹出
相当于把整个数组都放进堆里排序，再把前面多于k个的元素砍掉，剩下的第一个就是第k大的了

int findKthLargest(vector<int>& nums, int k) {
    priority_queue<int, vector<int>, greater<int>> q;
    for (int i = 0; i < nums.size(); ++i) {
        q.push(nums[i]);
        if (q.size() > k) q.pop();
    }
    return q.top();
}

method 2: 堆排序

void heapify(vector<int>& nums, int i, int n) {
    int l = 2 * i + 1, r = 2 * i + 2;
    int max = i;
    if (l < n && nums[l] > nums[max]) max = l;
    if (r < n && nums[r] > nums[max]) max = r;
    if (max != i) {
        swap(nums[i], nums[max]);
        heapify(nums, max, n);
    }
}
void buildHeap(vector<int>& nums, int n) {
    for (int i = n / 2; i >= 0; i--) {
        heapify(nums, i, n);
    }
}
int findKthLargest(vector<int>& nums, int k) {
    int n = nums.size();
    buildHeap(nums, n);
    for (int i = n - 1; i >= n - k + 1; i--) {
        swap(nums[0], nums[i]);
        heapify(nums, 0, i);
    }
    return nums[0];
}

method 3: 快速排序

第k大元素，也就是从小到大排序的倒数第k个元素，也就是index=nums.size()-k位置的元素，只要保证index左边的元素都比它小，右边的元素都比它大就行，至于左边和右边的元素是否有序不重要

partition选择一个主轴元素pivot，然后将小于pivot的元素放在其左边，大于它的元素放在它的右边，返回pivot的下标

quickSelect根据partition返回的下标判断，等于index的话就直接返回，如果小于index，说明index的元素还在右边，递归右边继续查找，否则递归左边查找

int quickSelect(vector<int>& nums, int l, int r, int index) {
    int ret = partition(nums, l, r);
    if (ret == index) {
        return nums[ret];
    }
    else if (ret < index) { // 小于index往右边找
        return quickSelect(nums, ret + 1, r, index);
    }   // 大于往左边找
    else return quickSelect(nums, l, ret - 1, index);
}
int partition(vector<int>& nums, int l, int r) {
    int pivot = nums[r];    // 直接以最右边的元素为主轴元素
    int i = l;
    for (int j = l; j <= r; j++) {
        if (nums[j] < pivot) {
            swap(nums[i], nums[j]);
            i++;
        }
    }   // [l,i-1]都比pivot小，[i,r-1]都比pivot大
    swap(nums[i], nums[r]); // 交换之后，[l,i-1]都比nums[i]小，[i+1,r]都比nums[i]大
    return i;
}
int findKthLargest(vector<int>& nums, int k) {
    int n = nums.size();
    return quickSelect(nums, 0, n - 1, n - k);
}