必会的QuickSelct算法 | partition函数，O(n)时间求第k大问题

问题描述

给一个未排序数组，要求输出第k大（小）的值。

算法分析

算法基本思想就是Quick Sort中的partition函数，我们将数组元素分成左右两部分，左边部分一定不大于右边部分（这里对于求第k小而言），根据我们两边划分的大小（或者说分界点的位置），继续递归地进行划分，直到达到我们需要的分界点处。

实际上我们可以直接使用STL中提供的partition函数，甚至我们可以直接使用nth_element函数解决第k大问题。

平均/最好时间复杂度 $O(n)$ ，最坏时间复杂度 $O(n^2)$

时间复杂度主要取决于选取的pivot （下面代码固定使用arr[r]能否将问题尽可能的缩小，最坏情况下问题大小只能减小1），可以随机选取pivot来使得我们的算法尽可能趋于平均复杂度（随机选取一个索引 $i \in [l, r]$ ，然后swap(arr[i], arr[r])）

示例代码

更详细的源代码在: github

#include <iostream>
#include <vector>

using namespace std;

/**
 * find the k-th (smallest) element of arr (k in [0..arr.size()-1])
 * e.g. the 3-th (smallest) element of {5,4,3,2,1,0} 3
 */
int quickSelect(vector<int> &arr, int l, int r, int k) {
  if (l == r) return arr[l];
  int ll = l, pivot = arr[r];
  for (int i = l; i < r; ++i) {
    if (arr[i] < pivot) {
      swap(arr[i], arr[ll++]);
    }
  }
  swap(arr[ll], arr[r]);
  // arr is like: ----(arr[ll]/pivot)++++, smaller ones on the left
  if (ll == k)
    return arr[k];
  else
    return ll < k ? quickSelect(arr, ll + 1, r, k)
                  : quickSelect(arr, l, ll - 1, k);
}

int main(int argc, char const *argv[]) {
  vector<int> arr{3, 2, 4, 5, 1};  // {1, 2, 3, 4, 5}
  int k = 3;
  cout << quickSelect(arr, 0, arr.size() - 1, k - 1) << endl;  // 3
  return 0;
}

扩展

partition的另一种方式
上面的示例代码采用的是Lomuto partition scheme, 还可以采用Hoare partition scheme，如下：(这个代码直接翻译的wiki的伪代码，个人觉得比较丑，参考中的csdn那篇给的实现比较好，最好的则是使用three way partition)

int quickSelect2(vector<int> &arr, int l, int r, int k) {
  if (l == r) return arr[l];
  int pivot = arr[(l + r) / 2];
  int ll = l - 1, rr = r + 1;
  while (true) {
    do
      ++ll;
    while (arr[ll] < pivot);
    do
      --rr;
    while (arr[rr] > pivot);
    if (ll >= rr) break;
    swap(arr[ll], arr[rr]);
  }
  // arr is like: ----++++, but the pivot index is unknown
  if (ll <= k)
    return quickSelect2(arr, ll, r, k);
  else
    return quickSelect2(arr, l, ll-1, k);
}

两种方式比较而言，Lomuto的更加简单，Hoare的效率上通常更高（不过pivot的位置不一定在划分处，参考中有保证分割点处为pivot的算法）

three-way-partition，简单而且可以分成---[=]=[=]+++三部分（[]表示这里的索引可以知道）

参考

wiki: Quickselect
wiki: Quicksort -- 两种partition方法
csdn: 划分(partition)算法 -- 类似Hoare的两端向中间扫描，但提前备份了pivot，保证分割点处为pivot，文中还有three way partition算法
wiki: Dutch national flag problem -- three way partition

最后编辑于：2020.08.17 18:59:19