【算法】快速选择

LeetCode 215. 数组中的第K个最大元素

给定整数数组 nums 和整数 k,请返回数组中第 k 个最大的元素。
请注意,你需要找的是数组排序后的第 k 个最大的元素,而不是第 k 个不同的元素。

求第K大元素、求第K小元素、求前K大元素、求前K小元素等问题都属于选择问题,即在某个序列中快速地选择某一(些)特定大小序位的元素。解决这些问题的一个直观办法是先进行排序,然后便可立即得到所要的结果,时间复杂度为O(nlogn),其实这个复杂度已经不慢了,但它并不是最优的。

为了排序某个数组,我们知道基于比较的排序算法,其理论最优时间复杂度就是O(nlogn),但在选择问题里,我们只关心特定顺序位置的元素,而不关心其他元素是否有序,而排序显然做了过多的事情。因此,是有空间去实现更快的时间复杂度的,本文将介绍的快速选择算法,具有O(n)的时间复杂度。

快速选择和快速排序名字类似,其实并不是巧合,因为快速选择的思想就来自于快速排序。在快速排序中,我们选择一个枢轴(Pivot),然后将数组分为三个部分:小于枢轴、大于枢轴、等于枢轴,以此实现排序。

注意到一个显然的事实:在每一次划分(Partition)过程中,我们可以确定至少一个元素在排序后的位置,即枢轴或者等于枢轴的元素在排序后的位置是确定的。如何确定呢?只需要看小于枢轴的元素有多少个就行了。

为了方便论述,假设我们现在要找第K小元素。根据这个思路,如果运气很好,选定的枢轴恰好是第K小元素,那么选择问题已经解决。当然实际情况下,枢轴不一定是第K小元素,但我们总是可以放弃掉一部分不可能是答案的元素。

设一次划分(Partition)后小于枢轴的元素集合为A,等于枢轴的元素集合为B,大于枢轴的元素集合为C,按如下减而治之

  • |A|<k
    • |A|+|B| \ge k,说明第K小元素等于枢轴,问题解决。
    • |A|+|B|<k,说明第K小元素在B中,对C递归,K减去|A|+|B|
  • |A| \ge k,说明第K小元素在A中,对A递归,K不变。

注意与快速排序有所不同,每一次划分后,我们只需要递归A或B中的一个集合,而不是两个。那么这个方法的效率就取决于如何划分数组了,确切地说,是如何确定一个枢轴,使得划分尽量均匀。在最坏的情况下,每次只排除一个元素的话,时间复杂度为O(n^2)

一种常用的方法是采用随机化确定枢轴,在这种情况下,可以证明时间复杂度的期望为O(n)。下面给出一种实现,该实现不采用随机化,却总能保证其时间复杂度为O(n)。该实现的步骤如下:

  • 步骤一:将数组分为每5个一组(不足5个的一组直接丢弃),取每一组的中值元素构成集合T
  • 步骤二:取集合T的中值元素,记为m
  • 步骤三:取m为枢轴,划分数组为三部分A,B,C,按照上述减而治之的思想递归

在三个步骤中,前两个步骤就是为了确定枢轴,这样的取法是有好处的。数组按每5个一组,组数为\lfloor \frac{n}{5} \rfloor。以一共5个组为例,在下图中,适当调整顺序,使组内元素从左往右、从小到大排序,且中值元素集合从上到下、从小到大排序。

则图中红色圆圈代表m,黄色区域中的元素必定大于等于m,绿色区域必定小于等于m。在黄色区域和绿色区域之外的元素,与m的大小关系并不确定,但已经足够让我们对两个区域的元素数量作一个估计了。

在一般情况下,共有\lfloor \frac{n}{5} \rfloor组,其中有一半的组的后三个元素均在黄色区域中,即大于等于枢纽的元素数量至少

    \[3\lceil \lfloor \frac{n}{5} \rfloor /2 \rceil \ge \frac{3}{2}\lfloor \frac{n}{5} \rfloor\]

那么小于等于枢纽的元素数量至多

    \[n - \frac{3}{2}\lfloor \frac{n}{5} \rfloor \le n - \frac{3}{2} ( \frac{n - 4}{5} ) = 0.7n + 1.2\]

同理,可以得到大于等于枢纽的元素数量也至多为0.7n + 1.2

即通过步骤一和步骤二方法确定的枢纽m,可以保证小于等于枢纽的元素数量和大于等于枢纽的元素数量具有一个上界,这个上界将在步骤三划分时保证算法的性能。

并且注意到步骤二取中值元素集合的中值,本质上也是一种选择问题,即K = \lceil \frac{n}{2} \rceil的情形,因而也可以通过递归实现。

接下来计算该算法的时间复杂度f(n)

  • 步骤一:将数组分为5个一组,取每组的中值元素构成集合。一共\lfloor \frac{n}{5} \rfloor组。对于一个长度为5的集合来说,只需常数时间便可得到其中值元素,故步骤一的时间复杂度为O(n)
  • 步骤二:取中值元素集合的中值,使用递归。该集合的长度为\lfloor \frac{n}{5} \rfloor,故步骤二的时间复杂度为f(\lfloor \frac{n}{5} \rfloor)
  • 步骤三:划分数组为三部分A,B,C,划分需要O(n)的时间。每次划分后只需递归A,C中的一个集合,而这两个集合的元素至多只有0.7n + 1.2个。当n \ge 38时,0.7n + 1.2 \le \lfloor \frac{3n}{4} \rfloor,即每一次划分后要递归的元素数量以\frac{3}{4}的比例递减。故步骤三的时间复杂度为O(n)+f(\lfloor \frac{3n}{4} \rfloor)

综上:n \ge 38, f(n) = O(n) + f(\lfloor \frac{n}{5} \rfloor) + f(\lfloor \frac{3n}{4} \rfloor),可以解得f(n) = O(n)

JAVA实现代码:

//快速选择数组前n个元素中的第k小元素
public static int Select(int[] a, int n, int k)
{
    //规模较小时直接排序
    if(n <= 38)
    {
        Arrays.sort(a, 0, n);
        return a[k - 1];
    }

    //取中值元素的中值为枢纽
    int[] ms = new int[n / 5];
    for(int i = 0; i < ms.length; i++)
    {
        ms[i] = Mid5(a, i);
    }
    int m = Select(ms, ms.length, (int)Math.ceil(ms.length / 2.0));

    //划分为三个集合
    int[] smaller = new int[3 * n / 4];
    int[] same  = new int[3 * n / 4];
    int[] larger = new int[3 * n / 4];
    int i = 0, j = 0, h = 0;
    for(int t = 0; t < n; t++)
    {
        if(a[t] < m) smaller[i++] = a[t];
        else if(a[t] == m) same[j++] = a[t];
        else larger[h++] = a[t];
    }

    //减而治之
    if(i >= k) return Select(smaller, i, k);
    else if(i + j >= k) return m;
    else return Select(larger, h, k - i - j);
}


//数组按5个一组分组后,取第i组的中值
private static int Mid5(int[] a, int i)
{
    int k = i * 5;
    //排序前四个元素
    if(a[k + 0] > a[k + 2]) Swap(a, k + 0, k + 2);
    if(a[k + 1] > a[k + 3]) Swap(a, k + 1, k + 3);
    if(a[k + 0] > a[k + 1]) Swap(a, k + 0, k + 1);
    if(a[k + 2] > a[k + 3]) Swap(a, k + 2, k + 3);
    if(a[k + 1] > a[k + 2]) Swap(a, k + 1, k + 2);

    //取得中值
    if(a[k + 4] > a[k + 2]) return a[k + 2];
    else if(a[k + 4] > a[k + 1]) return a[k + 4];
    else return a[k + 1];
}

private static void Swap(int[] a, int i, int j)
{
    int t = a[i];
    a[i] = a[j];
    a[j] = t;
}