前 K 个高频元素怎样找？哈希计数加堆和桶排序各是什么复杂度？

简答

LeetCode 347 给一个整数数组 nums 和整数 k，要返回出现次数最多的 k 个元素，顺序不限，题目保证答案唯一。比如 [1,1,1,2,2,3]、k = 2，1 出现 3 次、2 出现 2 次，答案是 [1,2]。进阶要求时间复杂度优于 O(n log n)，也就是不能把计数结果整个排序。

第一步都一样：用 Map 遍历一遍，记下每个数出现了几次，O(n)。难点在第二步，从计数结果里挑出次数最多的 k 个。

一种做法是小顶堆，和找第 K 大元素的思路相同，只是比较的是出现次数。堆里最多放 k 个，堆顶是次数最少的那个，来一个就放进去，超过 k 个就弹出堆顶。最多有 n 种不同的数，每次进出堆 O(log k)，总共 O(n log k)。另一种是桶排序：出现次数最多是 n，所以准备下标 0 到 n 的桶，把每个数放进下标等于它出现次数的桶里，再从下标 n 往下收集，凑够 k 个就停，总共 O(n)。

要注意两点。桶排序要额外开 n+1 个桶，数组很长但不同的数很少时会浪费空间，这时堆的做法更省。另外题目保证答案唯一，不需要处理第 k 名有几个数并列的情况，实际业务里如果有并列，要先定好怎样取舍。

两种写法

countFrequency 两种写法共用。MinHeap 和《数组中第 K 大的元素怎样找？小顶堆和快速选择的复杂度有什么区别？》里的写法相同，只是多了一个比较函数，因为这里堆里放的是 [数字, 次数]，要按次数比较。

JavaScript

function countFrequency(nums) {
  const counts = new Map();
  for (const num of nums) counts.set(num, (counts.get(num) ?? 0) + 1);
  return counts;
}

function topKByBucket(nums, k) {
  const counts = countFrequency(nums);
  // 出现次数最多是 nums.length，所以下标 0 到 n 的桶就够用。
  const buckets = Array.from({ length: nums.length + 1 }, () => []);
  for (const [num, count] of counts) buckets[count].push(num);

  const result = [];
  for (let count = nums.length; count > 0 && result.length < k; count -= 1) {
    for (const num of buckets[count]) {
      result.push(num);
      if (result.length === k) break;
    }
  }
  return result;
}

class MinHeap {
  constructor(compare) {
    this.items = [];
    this.compare = compare;
  }
  get size() {
    return this.items.length;
  }
  push(value) {
    const a = this.items;
    a.push(value);
    let i = a.length - 1;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (this.compare(a[parent], a[i]) <= 0) break;
      [a[parent], a[i]] = [a[i], a[parent]];
      i = parent;
    }
  }
  pop() {
    const a = this.items;
    const top = a[0];
    const last = a.pop();
    if (a.length) {
      a[0] = last;
      let i = 0;
      while (true) {
        let smallest = i;
        for (const child of [2 * i + 1, 2 * i + 2]) {
          if (child < a.length && this.compare(a[child], a[smallest]) < 0) smallest = child;
        }
        if (smallest === i) break;
        [a[smallest], a[i]] = [a[i], a[smallest]];
        i = smallest;
      }
    }
    return top;
  }
}

function topKByHeap(nums, k) {
  const counts = countFrequency(nums);
  // 堆里放 [数字, 次数]，按次数比较，堆顶是目前 k 个里次数最少的。
  const heap = new MinHeap((x, y) => x[1] - y[1]);
  for (const entry of counts) {
    heap.push(entry);
    if (heap.size > k) heap.pop();
  }
  return heap.items.map(([num]) => num);
}

const sort = (list) => [...list].sort((x, y) => x - y).join(',');
console.log(sort(topKByBucket([1, 1, 1, 2, 2, 3], 2))); // 1,2
console.log(sort(topKByHeap([1, 1, 1, 2, 2, 3], 2))); // 1,2
console.log(sort(topKByBucket([1, 2, 1, 2, 1, 2, 3, 1, 3, 2], 2))); // 1,2
console.log(sort(topKByHeap([4, 4, 5, 6, 6, 6, 7], 1))); // 6

题目允许任意顺序返回，所以打印前先排个序再比较。前三行用的是 LeetCode 题面里的例子，最后一行是额外补的：6 出现 3 次，比 4 的 2 次多，k = 1 时只返回 6。

桶排序为什么能做到 O(n)

普通排序要比较大小，所以至少 O(n log n)。桶排序不比较，而是利用了一个事实：出现次数是 1 到 n 之间的整数。既然范围已知，就可以直接拿次数当数组下标，把数放进对应的桶，这一步只要遍历一次计数结果。

收集时从下标 n 往下走，最多走 n+1 个桶。每个数只会被放进一个桶、最多被取出一次，所以收集也是 O(n)。三步加起来，计数 O(n)、分桶 O(n)、收集 O(n)，总共 O(n)。代价是空间，桶数组的长度是 n+1，不管有多少种不同的数都要开这么多。

堆为什么是小顶堆

要的是次数最多的 k 个，用的却是小顶堆，这一点容易绕。堆在这里像一个只有 k 个名额的排行榜，堆顶是榜上次数最少的那个，也就是“门槛”。新来的数如果次数比门槛高，放进去以后堆就超过 k 个了，弹出的是堆顶，也就是原来的门槛，新数就留下了。如果新来的数次数还不如门槛，它放进去以后自己就到了堆顶，马上又被弹出去。

所以遍历完，留在堆里的就是次数最多的 k 个。堆的大小始终不超过 k，每次进出 O(log k)，设不同的数有 m 种，总共 O(n + m log k)，m 最多是 n，所以写成 O(n log k)。堆的原理见《前端任务怎样按优先级排队？二叉堆为什么不需要每次完整排序？》。

常见追问

直接把计数结果排序可以吗？

可以，[...counts].sort((x, y) => y[1] - x[1]).slice(0, k) 写起来最短。但排序是 O(m log m)，m 是不同数的个数，最坏等于 n，就是 O(n log n)，不满足题目的进阶要求。面试时可以先说这个写法，再给出堆或桶排序作为优化。

数据是不断到来的流，怎么求当前的前 K 个？

桶排序需要知道全部数据，不适合流。可以一直维护计数 Map，需要结果时再用大小为 k 的小顶堆扫一遍计数，每次查询 O(m log k)。