使用最大堆解决“查找数组中第 K 个最大数”问题的时间复杂度是多少？

Question

"Find the K-th largest number in the array" problem:

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

I know that this problem can be soleved with quick select and min heap algorithms. However, what I'm focusing here is using max heap . The steps is the following:

1. Build max heap form the whole given array, inplace.
2. Iterate k times. In each iteration, Take and remove the heap top.
3. The last heap top taken at the k-th iteration is the result.

Below is the code in cpp:

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;
}

What is the time complextity of this method? I'm using bottom-up approach in the first step, this step should be O(n). However, the while loop makes me confused. The loop execute k times, each loop will call heapify_down which is O(log(n)) comlexity. So is the overall complexity O(n + k * log(n)) = O(max(n, k * log(n))) ? Correct me if I'm wrong.

Answer 1

Building a heap, if you do it right, is O(n) and removing the top is O(log n) and with k being unbound it could be n .

So overall it is O(n + k * log n) == O(n + n * log n) == O(n log n) .

Don't you have some more information, like k < log n or something?

---

When building your heap don't use heapify_down , because that results in O(n * log n). Instead first ignore the last element if the heap has odd size. Then start at i = size / 2 - 1 and swap i , 2 * i and 2 * i + 1 around so the largest is at i . Repeat till i-- = 0 . If the size was odd add the last element to the heap.