从一道简单算法题理解快速排序的 partition 操作

Posted 2021-05-02 五分钟学算法

点击蓝色“五分钟学算法”关注我哟

加个“星标”，天天中午 12:15，一起学算法

作者 | P.yh

来源 | 五分钟学算法

题目来源于 LeetCode 上第 75 号问题：颜色分类。题目难度为 Medium，目前通过率为 51.8% 。

题目描述

给定一个包含红色、白色和蓝色，一共 n 个元素的数组，原地对它们进行排序，使得相同颜色的元素相邻，并按照红色、白色、蓝色顺序排列。

此题中，我们使用整数 0、 1 和 2 分别表示红色、白色和蓝色。

注意:
不能使用代码库中的排序函数来解决这道题。

示例:

输入: [2,0,2,1,1,0]
输出: [0,0,1,1,2,2]

进阶：

• 一个直观的解决方案是使用计数排序的两趟扫描算法。
  首先，迭代计算出0、1 和 2 元素的个数，然后按照0、1、2的排序，重写当前数组。

• 你能想出一个仅使用常数空间的一趟扫描算法吗？

Follow up:

如果是 K 种颜色的话，该如何处理

题目解析

给定一个输入的数组，数组当中只有 0, 1, 2 这三种元素，让你对这个数组进行排序。我们先不急着去想最优解，来看看如果不加限制地看这道题，会有哪些解法：

• 利用归并排序，时间 O(nlogn)，空间 O(n)

• 利用快速排序，时间 O(nlogn)，空间 O(1)

• 利用计数排序，时间 O(n)，空间 O(1)

三种排序算法，显然计数排序会更优，因为这里只有 3 种元素，因此计数排序的空间复杂度也是常数级别的。

但是这道题最后问你的是能不能仅仅使用 One-Pass 来完成题目，One-Pass 的意思是仅仅只有一次 for 循环遍历，带着这个条件再来看这道题是不是会比较没有想法？

思路可以从 3 种颜色这里作为突破口，3 种颜色意味着排序好的数组，存在 3 个区域，每个区域存放的元素的值都是一样的：

[0...0,1...1,2...2]

我们可以想到用两个指针去维护其中的 2 个区域：

[0,...0,1...1,2...2]
 ------i     j-----

思路大概就有了，3 根指针，第一根维护 0 的区域，第二根维护 2 的区域，另外一根从头遍历数组，遇到 0 就和第一个指针指向的元素交换，遇到 2 就和第二个指针指向的元素交换，当遍历的指针和第二根指针相遇了就结束遍历。

这里有一个小细节就是，遍历的那根指针在交换完元素之后需不需要马上往前移动？

如果是和维护 0 的那根指针交换的话，因为遍历的这根指针已经遍历过这之前的所有元素了，因此交换完可以马上往前移动一个位置，但是如果是和维护 2 的那根指针交换的话，遍历的指针没有遍历过从那边交换过来的元素，交换过来的元素有可能是 0，有可能是 2，因此不能马上往前移动。

LeetCode 的上面这道题我们算是解决了，但是如果说这里不再是 3 种颜色，而是 K 种颜色，该如何处理呢？

如果是这种情况，像上面这种指针的做法就行不通了，你可能会想到那就直接排序吧，没错，排序的思路是对的，一般的快速排序，平均时间复杂度是 O(nlogn)，那能不能让他变得更快些，计数排序的话可以做到 O(n) 的时间复杂度，但是空间复杂度就会是 O(K)，如果这里还是要求你用常数级别的空间复杂度，该如何解决？

这里有一个点可能不太容易想到，平时我们想到快速排序，一般都知道，它的做法其实是利用分治的思想，把输入数组进行分割，对于这道题，需要换一种思路，就是我们基于颜色对数组进行分割，在分割数组的同时，我们也在分割颜色，这种做法可以把时间复杂度变成 O(nlogK)，因为颜色的数目肯定是小于元素的数目的，因此这个方法优于 O(nlogn)，具体可以参考下面的代码。

参考代码（一）：Sort Color

public void sortColors(int[] nums) {
    if (nums == null || nums.length == 0) {
        return;
    }

    int pointer0 = 0, pointerTraverse = 0, pointer2 = nums.length - 1;

    while (pointer2 >= pointerTraverse) {
        if (nums[pointerTraverse] == 0) { // 和维护 0 的指针交换元素，遍历指针往前移动
            swap(nums, pointer0++, pointerTraverse++);
        } else if (nums[pointerTraverse] == 2) {  // 和维护 2 的指针交换元素，遍历指针暂时不往前移动
            swap(nums, pointer2--, pointerTraverse);
        } else {
            pointerTraverse++;
        }
    }
}

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

参考代码（二）：Follow Up

public void sortColors2(int[] colors, int k) {
    if (colors == null || colors.length == 0 || colors.length <= k) {
        return;
    }

    quickSort(colors, 0, colors.length - 1, 1, k);
}

private void quickSort(int[] colors, 
                       int start, 
                       int end, 
                       int startColor, 
                       int endColor) {
    if (startColor >= endColor || start >= end) {
        return;
    }

    // 对颜色进行分割，并且每次都等分
    // 并利用中间的颜色作为数组的切分元素
    int midColor = (startColor + endColor) / 2;

    // 快速排序的思想，只不过这里 pivot 元素变成了上面选择的颜色
    int l = start, r = end;
    while (l <= r) {
        while (l <= r && colors[l] <= midColor) {
            l++;
        }

        while (l <= r && colors[r] > midColor) {
            r--;
        }

        if (l <= r) {
            swap(colors, l++, r--);
        }
    }

    // 同时基于数组和颜色进行分治
    quickSort(colors, start, r, startColor, midColor);
    quickSort(colors, l, end, midColor + 1, endColor);
}

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

如果这里理解不透，可以看之前的文章的动画进行理解 

这篇文章也可以加深你对 partition 操作的理解：

---

以上，便是今日分享，觉得不错，还请点个在看，谢谢～

推荐阅读：

啦，欢迎点击阅读原文进行访问~