75. 颜色分类

难度中等

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

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

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

示例:

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

进阶：

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

分析

循环不变量是这样定义的：

• 所有在子区间 [0, zero) 的元素都等于 0；
• 所有在子区间 [zero, i) 的元素都等于 1；
• 所有在子区间 [two, len - 1] 的元素都等于 2。

说明：设计循环不变量的原则是“不重不漏”。

1、len 是数组的长度；
2、变量 zero 是前两个子区间的分界点，一个是闭区间，另一个就必须是开区间；
3、变量 i 是循环变量，一般设置为开区间，表示 i 之前的元素是遍历过的；
4、two 是另一个分界线，我设计成闭区间。

于是编码要解决以下三个问题：

变量初始化应该如何定义。
在遍历的时候，是先加减还是先交换。
什么时候循环终止。

处理这三个问题，完全看循环不变量的定义。

• 编码的时候，zero 和 two 初始化的值就应该保证上面的三个子区间全为空。
• 在遍历的过程中，“索引先加减再交换”、还是“先交换再加减”就看初始化的时候变量在哪里。
• 退出循环的条件也看上面定义的循环不变量，在 i == two 成立的时候，上面的三个子区间就正好“不重不漏”地覆盖了整个数组，并且给出的性质成立，题目的任务也就完成了。

以下给出不同的写法，循环不变量的定义写在了注释中。使用“循环不变量”编码是为了便于我们处理细节，也方便别人看懂你的算法。

参考代码 1：

import java.util.Arrays;


public class Solution {

    public void sortColors(int[] nums) {
        int len = nums.length;
        if (len < 2) {
            return;
        }

        // all in [0, zero) = 0
        // all in [zero, i) = 1
        // all in [two, len - 1] = 2
        
        // 循环终止条件是 i == two，那么循环可以继续的条件是 i < two
        // 为了保证初始化的时候 [0, zero) 为空，设置 zero = 0，
        // 所以下面遍历到 0 的时候，先交换，再加
        int zero = 0;

        // 为了保证初始化的时候 [two, len - 1] 为空，设置 two = len
        // 所以下面遍历到 2 的时候，先减，再交换
        int two = len;
        int i = 0;
        // 当 i == two 上面的三个子区间正好覆盖了全部数组
        // 因此，循环可以继续的条件是 i < two
        while (i < two) {
            if (nums[i] == 0) {
                swap(nums, i, zero);
                zero++;
                i++;
            } else if (nums[i] == 1) {
                i++;
            } else {
                two--;
                swap(nums, i, two);
            }
        }
    }

    private void swap(int[] nums, int index1, int index2) {
        int temp = nums[index1];
        nums[index1] = nums[index2];
        nums[index2] = temp;
    }
}

参考代码 2：

public class Solution {

    public void sortColors(int[] nums) {
        int len = nums.length;
        if (len < 2) {
            return;
        }
        // all in [0, zero] = 0
        // all in (zero, i) = 1
        // all in (two, len - 1] = 2

        // 为了保证初始化的时候 [0, zero] 为空，设置 zero = -1，
        // 所以下面遍历到 0 的时候，先加，再交换
        int zero = -1;

        // 为了保证初始化的时候 (two, len - 1] 为空，设置 two = len - 1
        // 所以下面遍历到 2 的时候，先交换，再减
        int two = len - 1;
        int i = 0;
        // 当 i == two 的时候，还有一个元素还没有看，
        // 因此，循环可以继续的条件是 i <= two
        while (i <= two) {
            if (nums[i] == 0) {
                zero++;
                swap(nums, i, zero);
                i++;
            } else if (nums[i] == 1) {
                i++;
            } else {
                swap(nums, i, two);
                two--;
            }
        }
    }

    private void swap(int[] nums, int index1, int index2) {
        int temp = nums[index1];
        nums[index1] = nums[index2];
        nums[index2] = temp;
    }
}

说明：这种做法是我在 Java 的 JDK 的源码中 Arrays.sort() 中学到的。

我的体会是：

• 编码者应该在代码中使用注释表达这段代码编写的算法思想，提醒自己也方便他人。
• 但是源代码中类似 ++k <= great 和 a[++left] >= a[left - 1] 这样的代码建议不要写，会给阅读者带来理解上的障碍，这种做法也是《阿里巴巴 Java 开发手册》中不推荐的，理由是：变量的值发生变化是一个很重要的逻辑，应该单独成为一行，否则不利于调试和以后定位问题，这种语法糖我个人认为应该禁止使用。