题目来源于 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 - 11, 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;
}

如果这里理解不透,可以看之前的文章的动画进行理解 每天一算:Sort Colors

作者:P.yh 来源:https://juejin.im/post/5da8a0bcf265da5b894a1d34