Top K問題的兩種解決思路

Top K問題在數據分析中很是廣泛的一個問題（在面試中也常常被問到），好比：

從20億個數字的文本中，找出最大的前100個。

解決Top K問題有兩種思路，

最直觀：小頂堆（大頂堆 -> 最小100個數）； 較高效：Quick Select算法。 LeetCode上有一個問題215. Kth Largest Element in an Array，相似於Top K問題。

1. 堆

小頂堆（min-heap）有個重要的性質——每一個結點的值均不大於其左右孩子結點的值，則堆頂元素即爲整個堆的最小值。JDK中PriorityQueue實現了數據結構堆，經過指定comparator字段來表示小頂堆或大頂堆，默認爲null，表示天然序（natural ordering）。

小頂堆解決Top K問題的思路：小頂堆維護當前掃描到的最大100個數，其後每一次的掃描到的元素，若大於堆頂，則入堆，而後刪除堆頂；依此往復，直至掃描完全部元素。Java實現第K大整數代碼以下：

public int findKthLargest(int[] nums, int k) {
  PriorityQueue<Integer> minQueue = new PriorityQueue<>(k);
  for (int num : nums) {
    if (minQueue.size() < k || num > minQueue.peek())
      minQueue.offer(num);
    if (minQueue.size() > k)
      minQueue.poll();
  }
  return minQueue.peek();
}

2. Quick Select

Quick Select [1]脫胎於快排（Quick Sort），兩個算法的做者都是Hoare，而且思想也很是接近：選取一個基準元素pivot，將數組切分（partition）爲兩個子數組，比pivot大的扔左子數組，比pivot小的扔右子數組，而後遞推地切分子數組。Quick Select不一樣於Quick Sort的是其沒有對每一個子數組作切分，而是對目標子數組作切分。其次，Quick Select與Quick Sort同樣，是一個不穩定的算法；pivot選取直接影響了算法的好壞，worst case下的時間複雜度達到了O(n2)。下面給出Quick Sort的Java實現：

public void quickSort(int arr[], int left, int right) {
  if (left >= right) return;
  int index = partition(arr, left, right);
  quickSort(arr, left, index - 1);
  quickSort(arr, index + 1, right);
}

// partition subarray a[left..right] so that a[left..j-1] >= a[j] >= a[j+1..right]
// and return index j
private int partition(int arr[], int left, int right) {
  int i = left, j = right + 1, pivot = arr[left];
  while (true) {
    while (i < right && arr[++i] > pivot)
      if (i == right) break;
    while (j > left && arr[--j] < pivot)
      if (j == left) break;
    if (i >= j) break;
    swap(arr, i, j);
  }
  swap(arr, left, j);  // swap pivot and a[j]
  return j;
}

private void swap(int[] arr, int i, int j) {
  int tmp = arr[i];
  arr[i] = arr[j];
  arr[j] = tmp;
}

Quick Select的目標是找出第k大元素，因此

若切分後的左子數組的長度 > k，則第k大元素必出如今左子數組中； 若切分後的左子數組的長度 = k-1，則第k大元素爲pivot； 若上述兩個條件均不知足，則第k大元素必出如今右子數組中。 Quick Select的Java實現以下：

public int findKthLargest(int[] nums, int k) {
  return quickSelect(nums, k, 0, nums.length - 1);
}

// quick select to find the kth-largest element
public int quickSelect(int[] arr, int k, int left, int right) {
  if (left == right) return arr[right];
  int index = partition(arr, left, right);
  if (index - left + 1 > k)
    return quickSelect(arr, k, left, index - 1);
  else if (index - left + 1 == k)
    return arr[index];
  else
    return quickSelect(arr, k - index + left - 1, index + 1, right);

}

上面給出的代碼都是求解第k大元素；若想要獲得Top K元素，僅須要將代碼作稍微的修改：好比，掃描完成後的小頂堆對應於Top K，Quick Select算法用中間變量保存Top K元素。

3. 參考資料

[1] Hoare, Charles Anthony Richard. "Algorithm 65: find." Communications of the ACM 4.7 (1961): 321-322. [2] James Aspnes, QuickSelect.