使用max heap求解数组第K大元素的时间复杂度是多少

inputs: [3,2,1,5,6,4], k = 2
outputs: 5

inputs: [3,2,3,1,2,4,5,5,6], k = 4
outputs: 4

基于大顶堆（max heap）的实现方案 #

目前该问题常见解法包括quick select（快速选择）、min heap（小顶堆），这里聚焦max heap的实现逻辑，步骤如下：

• 基于给定的整个数组原地构建max heap
• 迭代k次，每次取出并移除堆顶元素
• 第k次迭代取出的堆顶元素即为最终结果

C++实现代码 #

void swap(vector<int>& nums, int i, int j) {
    int tmp = nums[i];
    nums[i] = nums[j];
    nums[j] = tmp;
}

void heapify_down(vector<int>& nums, int parent_idx, int end_idx) {
    int left_idx = parent_idx * 2 + 1, right_idx = left_idx + 1;
    while (left_idx <= end_idx) {
        int largeset_idx = parent_idx;
        if (nums[left_idx] > nums[largeset_idx]) largeset_idx = left_idx;
        if (right_idx <= end_idx && nums[right_idx] > nums[largeset_idx]) largeset_idx = right_idx;
        if (largeset_idx != parent_idx) {
            swap(nums, parent_idx, largeset_idx);
            parent_idx = largeset_idx;
            left_idx = parent_idx * 2 + 1;
            right_idx = left_idx + 1;
        }
        else {
            return ;
        }
    }
}

void build_heap(vector<int>& nums) {
    for (int i = nums.size() - 1; i >= 0; i--) heapify_down(nums, i, nums.size() - 1);
}

int findKthLargest(vector<int>& nums, int k) {
    build_heap(nums);
    int cnt = 0, res, cur_end = nums.size() - 1;
    while (cnt != k) {
        res = nums[0];
        cnt += 1;
        swap(nums, 0, cur_end);
        cur_end -= 1;
        heapify_down(nums, 0, cur_end);
    }

    return res;
}

时间复杂度解答 #

你的推导完全正确，分两部分验证：

1. 自底向上建堆的时间复杂度确实为O(n)。堆是完全二叉树结构，自底向上建堆时从最下层非叶子节点开始做下沉调整，每层节点的调整成本和节点高度对应，累加所有节点的调整成本后总开销为线性级别，不存在O(n log n)的冗余。
2. 后续k次取堆顶、调整堆的循环部分，严格来说每次调整的堆大小逐次减1，第i次调整的堆大小为n-i，对应下沉操作的时间复杂度为O(log(n-i))。由于log(n-i) ≤ log n恒成立，这部分的时间上界为O(k log n)。

两部分开销相加，整体时间复杂度为O(n + k log n)，等价于O(max(n, k log n))，符合不同场景下的复杂度表现：

• 当k为常数（如k=1找数组最大值），k log n为常数级别，总复杂度由建堆的O(n)主导
• 当k和n同阶（如k=n做完整堆排序找最小值），k log n为O(n log n)级别，总复杂度由堆调整部分主导，和堆排序的时间复杂度一致
• 当k远小于n/log n时，总复杂度为O(n)；当k增长到k log n > n后，总复杂度随k增长逐步上升到O(n log n)

内容的提问来源于stack exchange，提问作者YQ.Wang