87 lines
2.6 KiB
JavaScript
87 lines
2.6 KiB
JavaScript
/*
|
||
topk 问题是一个经典的问题，我们维持一个大小为k的堆，如果要求最大的第三个数我们就维持小顶堆，反之维持大小为k的大顶堆
|
||
为什么？这有点反直觉，为什么求第k大的数要维持小顶堆，而不是大顶堆，这是因为第k大的数是数组中k个最大的数中最小的那一个，
|
||
用小顶堆来维持那么要求的第k个最大的数刚好在小顶堆的堆顶
|
||
*/
|
||
|
||
const numbers = [5, 0, 8, 2, 1, 4, 7, 2, 5];
|
||
|
||
/**
|
||
*
|
||
* @param {number[]} nums
|
||
* @param {number} k
|
||
*/
|
||
function topK(nums, k) {
|
||
const heap = Array(k).fill(-Infinity); // 定义一个大小为k数组用来维护堆
|
||
|
||
for (let i = 0; i < nums.length; i++) {
|
||
// 如果当前元素大于堆顶元素，替换堆顶元素然后下沉调整堆
|
||
if (nums[i] > heap[0]) {
|
||
heap[0] = nums[i];
|
||
let cur = 0; // cur指向下沉元素的下标,一开始就在堆顶
|
||
while (true) {
|
||
const left = cur * 2 + 1;
|
||
const right = cur * 2 + 2;
|
||
let smallest = cur;
|
||
|
||
if (left < k && heap[left] < heap[smallest]) smallest = left;
|
||
if (right < k && heap[right] < heap[smallest]) smallest = right;
|
||
|
||
// 如果smallest 和cur相等，表示当前元素来到了合适的位置cur，结束下沉
|
||
if (smallest === cur) break;
|
||
|
||
// 交换位置下沉
|
||
[heap[cur], heap[smallest]] = [heap[smallest], heap[cur]];
|
||
cur = smallest;
|
||
}
|
||
}
|
||
}
|
||
|
||
// 返回堆顶元素
|
||
return heap[0];
|
||
}
|
||
|
||
/**
|
||
* 使用小顶堆找出数组中第 K 大的数
|
||
* @param {number[]} nums
|
||
* @param {number} k
|
||
* @returns {number}
|
||
*/
|
||
function topKbest(nums, k) {
|
||
const heap = nums.slice(0, k); // 先取前 k 个元素构建小顶堆
|
||
|
||
// 下沉函数（维护堆性质）
|
||
function siftDown(heap, start, heapSize) {
|
||
let cur = start;
|
||
while (true) {
|
||
const left = 2 * cur + 1;
|
||
const right = 2 * cur + 2;
|
||
let smallest = cur;
|
||
|
||
if (left < heapSize && heap[left] < heap[smallest]) smallest = left;
|
||
if (right < heapSize && heap[right] < heap[smallest]) smallest = right;
|
||
|
||
if (smallest === cur) break;
|
||
|
||
[heap[cur], heap[smallest]] = [heap[smallest], heap[cur]];
|
||
cur = smallest;
|
||
}
|
||
}
|
||
|
||
// 构建小顶堆（自底向上建堆）
|
||
for (let i = Math.floor(k / 2) - 1; i >= 0; i--) {
|
||
siftDown(heap, i, k);
|
||
}
|
||
|
||
// 遍历剩余元素，维护一个大小为 k 的小顶堆
|
||
for (let i = k; i < nums.length; i++) {
|
||
if (nums[i] > heap[0]) {
|
||
heap[0] = nums[i];
|
||
siftDown(heap, 0, k);
|
||
}
|
||
}
|
||
|
||
// 堆顶即为第 k 大元素
|
||
return heap[0];
|
||
}
|