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

本是古典何须时尚 2023-05-31 02:38 30阅读 0赞

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

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

![640?wx\_fmt=jpeg][640_wx_fmt_jpeg]

作者 | P.yh

来源 | 五分钟学算法

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

### 题目描述 ###

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

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

**注意:**

**示例:**

输入: [2,0,2,1,1,0]
    输出: [0,0,1,1,2,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 操作的理解：

\---

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