1. 题目描述

输入 n 个整数,找出其中最小的 k 个数。例如输入 4、5、1、6、2、7、3、8 这 8 个数字,则最小的 4 个数字是 1、2、3、4。

2. 思路分析

这里创建一个容量为 k 的最大堆。遍历给定数据集合,每次和堆顶元素进行比较,如果小于堆顶元素,则弹出堆顶元素,然后将当前元素放入堆。

由于堆大小为 k,所以弹出、推入操作复杂度为:O(logK)。因为有 n 个,总体复杂度为:O(nLogK)。

对比快排 partition 的思路,这种思路优点如下:

  1. 不会变动原数组
  2. 适合处理海量数据,尤其对于不是一次性读取的数据

3. 代码实现

请先执行:yarn add heap 或者 npm install heap

代码如下;

const Heap = require("heap");

function compare(a, b) {
  if (a < b) {
    return 1;
  }
  if (a > b) {
    return -1;
  }
  return 0;
}

function getKthNumbers(nums = [], k) {
  if (k <= 0) {
    return null;
  }

  const heap = new Heap(compare);
  for (let num of nums) {
    if (heap.size() < k) {
      heap.push(num);
    } else {
      const top = heap.pop();
      if (num <= top) {
        heap.push(num);
      } else {
        heap.push(top);
      }
    }
  }

  return heap.toArray();
}

/**
 * 以下是测试代码
 */

console.log(getKthNumbers([4, 5, 1, 6, 2, 7, 3, 8], 4)); // output: [ 4, 3, 1, 2 ]
console.log(getKthNumbers([10, 2], 1)); // output: [ 2 ]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40