我是前端西瓜哥，今天來整下 TopK 算法。

TopK，即求數(shù)組的最小(或最大)的 k 個數(shù)，且不要求這些數(shù)要排序返回。

這是一個非常經(jīng)典的面試題。解法也是相當(dāng)?shù)亩?，能較好考察面試者的數(shù)據(jù)結(jié)構(gòu)與算法基礎(chǔ)。

求最小和最大的思路是一樣的，我們假設(shè)要求的是最小的 k 個數(shù)。

對應(yīng)的 LeetCode 題目地址有兩個：

• 《劍指 Offer(第 2 版)》第 40 題：https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof/
• 《程序員面試經(jīng)典(第 6 版)》17.14 題：https://leetcode-cn.com/problems/smallest-k-lcci/

排序

最簡單的方式是全排序，即對數(shù)組的所有元素都進(jìn)行升序排序，然后取前面的 k 個數(shù)。通常都會用編程語言自帶的排序 API，正確性有所保證。

function getLeastNumbers(arr: number[], k: number): number[] {
  return arr.sort((a, b) => a - b).slice(0, k);
};

實際開發(fā)中如果有這個需求，且數(shù)據(jù)量不大，用自帶的排序方法是最穩(wěn)妥的。

因為排序方法底層通常是快排這些時間復(fù)雜度優(yōu)秀的排序算法，所以全排序的時間復(fù)雜度是 O(n*log(n)。

在全排序的基礎(chǔ)上，其實可以做個優(yōu)化，做 局部排序。有些算法(冒泡和選擇排序)的排序過程中，第 i 次迭代，都會使得第 i 個元素置于最終排好序后所在的位置。

這里我們看看魔改選擇排序的實現(xiàn)：

function getLeastNumbers(arr: number[], k: number): number[] {
  let min: number;
  let minIdx: number;
  for (let i = 0; i < k; i++) { // 這里迭代了 k 次
    min = arr[i];
    minIdx = i;
    for (let j = i + 1; j < arr.length; j++) {
      if (arr[j] < min) {
        min = arr[j];
        minIdx = j;
      }
    }
    [arr[i], arr[minIdx]] = [arr[minIdx], arr[i]]; // 交換
  }
  return arr.slice(0, k);
};

只要我們將原來的 n 次迭代改為 k 次迭代，就能獲得一個只是前 k 個數(shù)組元素做了排序的數(shù)組。

局部排序在原來時間復(fù)雜度為 O(n^2) 的排序算法的基礎(chǔ)上，優(yōu)化為了 O(k*n)。

當(dāng) k 很小時，局部排序的效率要比全排序的高。

快排思想

快速排序的核心是 分治 和 分區(qū)。

限于篇幅，具體的快排原理和實現(xiàn)可以看我的這篇文章：經(jīng)典快速排序?qū)崿F(xiàn)

簡單來說，快排的實現(xiàn)是，每次取一個基準(zhǔn)值 pivot，將小于等于 pivot 的放到左區(qū)間，大于的放到右區(qū)間。然后針對這兩個區(qū)間繼續(xù)同樣操作，直到區(qū)間大小小于等于 1 為止，是自上而下的遞歸。

原來快排對兩個區(qū)間都要進(jìn)行遞歸，現(xiàn)在改造為選擇性地遞歸。

每一次分區(qū)(partition)后：

• 如果 pivot 所在的位置小于 k，遞歸就可以結(jié)束了，此時數(shù)組的前 k 個數(shù)組元素就是最小的 k 個元素;
• 如果 pivot 所在的位置在 k 的左側(cè)，說明 pivot 的左區(qū)間可以不用排序了，小于等于 pivot 的值都在左側(cè)。然后對右區(qū)間進(jìn)行遞歸;
• 如果 pivot 所在的位置在 k 的右側(cè)，則 pivot 的右區(qū)間不用考慮了，需要對左區(qū)間遞歸。

這里有一個非常重要的點：每次分區(qū)分后 pivot 所在索引的值就是整個數(shù)組排過序后應(yīng)該為的值。 這是我們可以不管其中一個區(qū)間的原因。

function getLeastNumbers(arr: number[], k: number): number[] {
  partition(arr, 0, arr.length - 1, k);
  return arr.slice(0, k);
};

function partition(arr: number[], lo: number, hi: number, k: number) {
  if (lo >= hi) return;
  const pivot = arr[hi];  // 這里可以改為隨機選擇 pivot
  let i = lo;
  for (let j = lo; j < hi; j++) {
    if (arr[j] <= pivot) {
      swap(arr, i, j);
      i++;
    }
  }
  swap(arr, i, hi);
  
  // 原本的快排的 partition 的最后是這兩行，現(xiàn)在改為現(xiàn)在的下面這五行
  // partition(arr, i + 1, hi, k);
  // partition(arr, lo, i - 1, k);
  if (i < k) {
    partition(arr, i + 1, hi, k);
  } else if (i > k) {
    partition(arr, lo, i - 1, k);
  }
}

function swap(arr: number[], i: number, j: number) {
  let tmp = arr[i];
  arr[i] = arr[j];
  arr[j] = tmp;
}

平均時間復(fù)雜度是 O(n)，最壞時間復(fù)雜度是 O(n^2)。不管怎樣總體上都比快排效率高，除非極端情況。

大頂堆

大頂堆是個完全二叉樹，特點是：任何一個節(jié)點的值都大于等于子樹任意一個節(jié)點的值。

我們創(chuàng)建一個大小為 k 的大頂堆。先放入 k 個元素。后面想放入新元素時，先和堆頂元素(堆的最大值)對比。如果小于堆的最大元素，說明這個堆頂元素不合格，不可能為 TopK 的一員，將其出堆，然后讓新元素入堆;否則新元素不入堆。

一直這樣操作直到遍歷完整個數(shù)組。最后這個堆就是我們要的 TopK。

JavaScript 沒有內(nèi)置堆或優(yōu)先隊列這些數(shù)據(jù)結(jié)構(gòu)，就暫且不實現(xiàn)了。

入堆的時間復(fù)雜度是 O(log k)，要入堆 n 次，所以總的時間復(fù)雜度是 O(n*log k)。

如果你需要動態(tài)維護 TopK，比如網(wǎng)站的每日排行榜，用大頂堆方案會更合適。

結(jié)尾

總的來說，快排思想的方案時間復(fù)雜度最低，大頂堆適合需要動態(tài)維護 TopK 的情況，而全排序則適合寫業(yè)務(wù)代碼且數(shù)據(jù)量不大的時候。