algorighm/heap/base/topk问题.js

132 lines
4.3 KiB
JavaScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

/*
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];
}
/*
若题目没有要求动态的维护这个topk,那么可以使用quick-select来实现O(n)的时间复杂度完成这个题目。
思路题目要求我们找出第k大的数那么当整个数组排序之后倒数第k个数不就是我们要求的吗比如[5,4,3,2,1] k = 2,那么排序后[1,2,3,4,5]中
倒数第二个数的下标就是nums.length - k = 5 - 2 = 3 nums[3]恰好就是我们要找的这个数,我们没有必要对整个数组排序,我们只要找到排序后这个
数正确的位置即可知道快速排序的话我们可以使用partition来确定pivot,pivot在partition之后就已经确认如果pivot的下标等于倒数第k个位置
那么pivot就是排序之后的倒数第k个数也就是topk
*/
function f2(nums, k) {
return quickSelect(nums, 0, nums.length - 1, nums.length - k);
}
function quickSelect(nums, left, right, n) {
const p = partition(nums, left, right);
if (p === n) return nums[n];
return p < n ? quickSelect(nums, p + 1, right, n) : quickSelect(nums, left, p - 1, n);
}
/**
* 分区函数返回pivotpivot 左边的数都不大于,右边的数都不小于它
* @param {*} nums 要操作的数组
* @param {*} left 开始位置,左边界
* @param {*} right 结束位置, 右边界
* @returns number 返回的piivot
*/
function partition(arr, left, right) {
const pivot = arr[right];
let p = left;
for (let i = left; i < right; i++) {
if (arr[i] <= pivot) {
// [arr[i], arr[p]] = [arr[p], arr[i]];
swap(arr, i, p);
p++;
}
}
// [arr[p], arr[right]] = [arr[right], arr[p]];
swap(arr, p, right);
return p;
}
function swap(nums, a, b) {
const tmp = nums[a];
nums[a] = nums[b];
nums[b] = tmp;
}