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

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

這是一個非常經典的面試題。解法也是相當的多，能較好考察面試者的數據結構與算法基礎。

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

對應的 LeetCode 題目地址有兩個：

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

排序

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

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

實際開發中如果有這個需求，且數據量不大，用自帶的排序方法是最穩妥的。

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

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

這里我們看看魔改選擇排序的實現：

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 個數組元素做了排序的數組。

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

當 k 很小時，局部排序的效率要比全排序的高。

快排思想

快速排序的核心是 分治 和 分區。

限于篇幅，具體的快排原理和實現可以看我的這篇文章：經典快速排序實現

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

原來快排對兩個區間都要進行遞歸，現在改造為選擇性地遞歸。

每一次分區(partition)后：

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

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

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 的最后是這兩行，現在改為現在的下面這五行
  // 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;
}

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

大頂堆

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

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

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

JavaScript 沒有內置堆或優先隊列這些數據結構，就暫且不實現了。

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

如果你需要動態維護 TopK，比如網站的每日排行榜，用大頂堆方案會更合適。

結尾

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