小根堆-前 K 个高频元素

347. 前 K 个高频元素

给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。

示例 1:

输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]

示例 2:

输入: nums = [1], k = 1
输出: [1]

提示：

• 1 <= nums.length <= 105
• k 的取值范围是 [1, 数组中不相同的元素的个数]
• 题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的

进阶： 你所设计算法的时间复杂度 必须 优于 O(n log n) ，其中 n **是数组大小。

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number[]}
 */
var topKFrequent = function(nums, k) {
	const map = new Map();
	for (const num of nums) {
		if (map.has(num)) {
			map.set(num, map.get(num) + 1);
		} else {
			map.set(num, 1);
		}
	}

	const heap = new Heap((a, b) => a[1] < b[1]);

	for (const [key, count] of map) {
		if (heap.getSize() < k) {
			heap.push([key, count])
		} else {
			const [, topCount] = heap.getTop();
			if (count > topCount) {
                heap.pop();
				heap.push([key, count]);
			}
		}
	}
	
	const result = [];
	const nodeList = heap.nodeList;
	for (const [key, ] of nodeList) {
		result.push(key);
	}

	return result;
};

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