本题其实是经典的荷兰国旗问题，最初由荷兰计算机科学家艾兹赫尔·迪克斯特拉（Edsger W. Dijkstra）提出，并以荷兰国旗的颜色（红、白、蓝）为类比。

如果你学习过「快速排序」，知道「三路快排」，这道题就非常简单，本题考查的知识点就是「三路快排」。我以前录过视频讲解了「快速排序」：地址，在第 6 节讲到了「三路快排」。

思路分析

容易想到的做法：

• 排序。可以使用编程语言提供的排序函数，一般认为是「快速排序」或者「归并排序」，排序以后即为所求。但不符合题目的「进阶」要求「一趟扫描」和「常数空间」；
• 使用「计数排序」。分别统计 0、1、2 的个数，再赋值回原数组。但不符合题目的「进阶」要求「一趟扫描」。

「快速排序」的 partition，其中有一种 partition 叫做「三路快排」，即通过一次 partition，把区间里的数根据基准元素 pivot 分成 3 个部分：

• 第 1 个部分：小于 pivot；
• 第 2 个部分：等于 pivot；
• 第 3 个部分：大于 pivot。

本题的解法和「三路快排」几乎是一样，在「一趟扫描」的过程中，使用变量 i 扫描，分别使用两个变量放在数组的头和尾，我们就分别命名为 zero 和 two。其中

• zero 是 0 和 1 的分界；
• two 是 2 和还未遍历到的数的分界。

变量的定义如下图所示（「参考代码 1」采用的定义方式）：

变量 i 看到 1 的时候直接 ++ ，看到 0 的时候交换到数组的前面，看到 2 的时候交换到数组的末尾。具体细节，可以先写出变量的定义，在编码就容易多了。

这里给出两版代码，其实是一样的，它们的区别仅仅在于 zero 指向 0 还是指向 1 ，two 指向 2 还是指向未看到的数。我们把 zero 和 two 的定义写成了区间的形式，写在注释中。

由于区间定义不同，初始化，循环过程中，退出循环的条件就有细微差别。

参考代码 1：

###java

public class Solution {

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

        // all in [0..zero) = 0
        // all in [zero..i) = 1
        // all in [two..n - 1] = 2
        // 初始化的时候，要满足上面 3 个区间全是空区间
        int zero = 0;
        int two = n;
        int i = 0;
        // i = two 的时候，[zero..i)、[two..len - 1] 已经接起来了，所以 while 里面是 i < two
        while (i < two) {
            if (nums[i] == 0) {
                // zero 指向 1，所以先交换
                swap(nums, i, zero);
                zero++;
                i++;
            } else if (nums[i] == 1) {
                // 直接划分到 1 所在的区间
                i++;
            } else {
                // [two..n - 1] = 2，two 指向 2，所以先 -- ，再交换
                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;
    }

}

复杂度分析：

• 时间复杂度：$O(n)$，这里 $n$ 是输入数组的长度；
• 空间复杂度：$O(1)$。

参考代码 2：

###java

public class Solution {

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

        // all in [0, zero] = 0
        // all in (zero, i) = 1
        // all in (two, n - 1] = 2
        int zero = -1;
        int two = n - 1;
        int i = 0;
        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;
    }

}

复杂度分析：（同「参考代码 1」）。