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

73 阅读1分钟

题目链接

给定整数数组 nums 和整数 k,请返回数组中第 k 个最大的元素。

请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

解法 小根堆

思路

手写小根堆其实核心就在 bubbleUpbubbleDown 这两个函数。

插入一个元素后,调用 bubbleUp ,用它和它的父节点比大小,把它“冒”上去到合适的位置。

删除堆顶(最小堆中是最小值)后,调用 bubbleDown, 把最后一个元素放到堆顶,然后和它的左右子节点比较,让它“沉”下去。

代码

class MinHeap {
    private heap: number[] = [];

    insert(val: number) {
        this.heap.push(val);
        this.bubbleUp();
    }

    pop(): number | undefined {
        const top = this.heap[0];
        const end = this.heap.pop();
        if (this.heap.length > 0 && end !== undefined) {
            this.heap[0] = end;
            this.bubbleDown();
        }
        return top;
    }

    peek(): number | undefined {
        return this.heap[0];
    }

    size(): number {
        return this.heap.length;
    }

    private bubbleUp() {
        let index = this.heap.length - 1;
        while (index > 0) {
            const parent = Math.floor((index - 1) / 2);
            if (this.heap[index] >= this.heap[parent]) break;
            [this.heap[index], this.heap[parent]] = [this.heap[parent], this.heap[index]];
            index = parent;
        }
    }

    private bubbleDown() {
        let index = 0;
        const length = this.heap.length;
        const element = this.heap[0];

        while (true) {
            let left = 2 * index + 1;
            let right = 2 * index + 2;
            let smallest = index;

            if (left < length && this.heap[left] < this.heap[smallest]) {
                smallest = left;
            }

            if (right < length && this.heap[right] < this.heap[smallest]) {
                smallest = right;
            }

            if (smallest === index) break;
            [this.heap[index], this.heap[smallest]] = [this.heap[smallest], this.heap[index]];
            index = smallest;
        }
    }
}

function findKthLargest(nums: number[], k: number): number {
    const minHeap = new MinHeap();
    for (let num of nums) {
        minHeap.insert(num);
        if (minHeap.size() > k) {
            minHeap.pop();
        }
    }
    return minHeap.peek();
};

时空复杂度

时间复杂度:O(n logk)

空间复杂度:O(k)