Code Top -- 04 数组中的第K个最大元素

数组中的第K个最大元素

Posted by OUC_LiuX on March 28, 2022

from CodeTop, Leetcode, 215. 数组中的第K个最大元素
用堆实现,时间空间消耗比快排低。

给定整数数组 nums 和整数 k,请返回数组中第 k 个最大的元素。
请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。

实用大根堆实现

重建一个大根堆,找第 k 大元素,则顺序删除 k-1 次堆顶元素,返回堆顶即可。

static const int N = 10010;
int heap[N], cap;

int findKthLargest(vector<int>& nums, int k){
    cap = 0;
    for (int x: nums){
        heap[++cap] = x;
        up(cap);
    }

    for (int i = 1; i < k; i++) {
        heap[1] = heap[cap--];
        down(1);
    }

    return heap[1];
}

void up(int u){
    if (u/2 && heap[u/2] < heap[u]){
        swap(heap[u], heap[u/2]);
        up(u/2);
    }
}

void down(int u){
    int t = u;
    if (2*u <= cap && heap[2*u] > heap[t])  t = 2*u;
    if (2*u+1 <= cap && heap[2*u+1] > heap[t])  t = 2*u+1;
    if (t > u){
        swap(heap[u], heap[t]);
        down(t);
    }
}