颜色分类

题目介绍

力扣75题：https://leetcode-cn.com/problems/sort-colors/
给定一个包含红色、白色和蓝色，一共 n 个元素的数组，原地对它们进行排序，使得相同颜色的元素相邻，并按照红色、白色、蓝色顺序排列。

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

分析

本题是经典的“荷兰国旗问题”，由计算机科学家 Edsger W. Dijkstra 首先提出。荷兰国旗是由红白蓝3种颜色的条纹拼接而成，如下图所示：

假设这样的条纹有多条，且各种颜色的数量不一，并且随机组成了一个新的图形，新的图形可能如下图所示，但是绝非只有这一种情况：

需求是：把这些条纹按照颜色排好，红色的在上半部分，白色的在中间部分，蓝色的在下半部分，我们把这类问题称作荷兰国旗问题。本题其实就是荷兰国旗问题的数学描述，它在本质上，其实就是就是一个有重复元素的排序问题。所以可以用排序算法来解决。当然，最简单的方式，就是直接调Java已经内置的排序方法：

public void sortColors(int[] nums) {
    Arrays.sort(nums);
}

时间复杂度为O(nlogn)。但显然这不是我们想要的，本题用到的排序算法应该自己实现，而且要根据本题的具体情况进行优化。

方法一：基于选择排序

如果用选择排序的思路，我们可以通过遍历数组，找到当前最小（或最大的数）。对于本题，因为只有0，1，2三个值，我们不需要对每个位置的“选择”都遍历一遍数组，而是最多遍历三次就够了：第一次遍历，把扫描到的0全部放到数组头部；第二次遍历，把所有1跟在后面；最后一次，把所有2跟在最后。事实上，最后对于2的扫描已经没有必要了：因为除了0和1，剩下的位置一定都是2。所以我们可以用两次扫描，实现这个算法。

代码如下：

public void sortColors(int[] nums) {
    int curr = 0;
    // 第一次遍历，将扫描到的0交换到数组头部
    for ( int i = 0; i < nums.length; i++){
        if ( nums[i] == 0 ){
            swap( nums, curr++, i ); 
        }
    }
    // 第二次遍历，将扫描到的1跟在后面
    for ( int i = 0; i < nums.length; i++){
        if ( nums[i] == 1 ){
            swap( nums, curr++, i ); 
        }
    }
}
public void swap( int[] nums, int i, int j ){
    int temp = nums[i];
    nums[i] = nums[j];
    nums[j] = temp;
}

复杂度分析

• 时间复杂度：O(n)，n为数组nums的长度。需要遍历两次数组。
• 空间复杂度：O(1)，只用到了常数个辅助变量。

方法二：基于计数排序

据题目中的提示，要排序的数组中，其实只有0，1，2三个值。
所以另一种思路是，我们可以直接统计出数组中 0,1,2 的个数，再根据它们的数量，重写整个数组。这其实就是计数排序的思路。

代码如下：

public class SortColors {
    public void sortColors(int[] nums) {
        int count0 = 0, count1 = 0, count2 = 0;
        // 遍历数组，统计0，1，2的个数
        for ( int num: nums ){
            if ( num == 0 )
                count0 ++;
            else if ( num == 1 )
                count1 ++;
            else
                count2 ++;
        }
        // 将0，1，2按个数依次填入数组
        for ( int i = 0; i < nums.length; i++ ){
            if ( i < count0 )
                nums[i] = 0;
            else if ( i < count0 + count1 )
                nums[i] = 1;
            else
                nums[i] = 2;
        }
    }
}

复杂度分析

• 时间复杂度：O(n)，n为数组nums的长度。需要遍历两次数组。
• 空间复杂度：O(1)，只用到了常数个辅助变量。

方法三：基于快速排序

前面的算法，尽管时间复杂度为O(n)，但都进行了两次遍历。能不能做一些优化，只进行一次遍历就解决问题呢？

一个思路是，使用双指针。所有的0移到数组头，所有2移到数组尾，1保持不变就可以了。这其实就是快速排序的思路。

代码如下：

public void sortColors(int[] nums) {
    int left = 0, right = nums.length - 1;
    int i = left;
    while ( left < right && i <= right ){
        while ( i <= right && nums[i] == 2 )
            swap( nums, i, right-- );
        if ( nums[i] == 0 )
            swap( nums, i, left++ );
        i++;
    }
}

复杂度分析

• 时间复杂度：O(n)，n为数组nums的长度。双指针法只虚遍历一次数组。
• 空间复杂度：O(1)，只用到了常数个辅助变量。