小根堆-数组中的第K个最大元素

160 阅读1分钟

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

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

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

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

 

示例 1:

输入: [3,2,1,5,6,4], k = 2
输出: 5

示例 2:

输入: [3,2,3,1,2,4,5,5,6], k = 4
输出: 4

 

提示:

  • 1 <= k <= nums.length <= 105
  • -104 <= nums[i] <= 104
/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var findKthLargest = function(nums, k) {
	const heap = new Heap((a, b) => a < b);

	for (const num of nums) {
		if (heap.getSize() < k) {
			heap.push(num);
		} else {
			if (num > heap.getTop()) {
				heap.pop();
				heap.push(num);
			}
		}
	}
	return heap.getTop();
};


class Heap {
	constructor(compare) {
		this.nodeList = [];
		this.compare = typeof compare === 'function' ? compare : this.defaultCompare;
	}

	defaultCompare(a, b) {
		return a > b;
	}

	getSize() {
		return this.nodeList.length;
	}

	getTop() {
		if (this.nodeList.length === 0) {
			return null;
		}

		return this.nodeList[0];
	}

	push(value) {
		this.nodeList.push(value);
		this.up(this.nodeList.length - 1);
	}

	up(index) {
		const { compare, parent, nodeList } = this;
		let curIndex = index;
		let parentIndex = parent(curIndex);
		while (curIndex > 0 && compare(nodeList[curIndex], nodeList[parentIndex])) {
			const temp = nodeList[curIndex];
			nodeList[curIndex] = nodeList[parentIndex];
			nodeList[parentIndex] = temp;
			curIndex = parentIndex;
			parentIndex = parent(curIndex);
		}
	}

	parent(index) {
		return index % 2 === 0 ? index / 2 - 1 : (index - 1) / 2;
	}

	pop() {
		if (this.nodeList.length === 0) {
			return null;
		}

		const top = this.nodeList[0];
		this.nodeList[0] = this.nodeList[this.nodeList.length - 1];
		this.nodeList.pop();
		this.down(0);
		return top;
	}

	down(index) {
		const { compare, left, right, nodeList } = this;
		let curIndex = index;
		while (left(curIndex) < nodeList.length) {
			let target = left(curIndex);
			if (right(curIndex) < nodeList.length && compare(nodeList[right(curIndex)], nodeList[left(curIndex)])) {
				target = right(curIndex);
			}

			if (compare(nodeList[curIndex], nodeList[target])) {
				break;
			}

			const temp = nodeList[curIndex];
			nodeList[curIndex] = nodeList[target];
			nodeList[target] = temp;
			curIndex = target;
		}
	}

	left(index) {
		return index * 2 + 1;
	}

	right(index) {
		return index * 2 + 2;
	}
}