题目描述

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

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

必须在不使用库内置的 sort 函数的情况下解决这个问题。

示例 1：

1
2
输入：nums = [2,0,2,1,1,0]
输出：[0,0,1,1,2,2]

算法思路

这个题目如果遍历两次的话非常好解决，首先第一遍遍历，把0换到数组最左边，使用一个指针p，划定0的范围，然后遍历数组，如果nums[i] = 0，就执行swap(nums[i], nums[pre]), pre++；第二遍遍历是一样的，这次从pre开始，再把1换到以pre起始的地方。python伪代码（方便点）大概是：

1
2
3
4
5
6
7
8
for i in range(0, n):
	if nums[i] == 0:
		swap(nums[pre], nums[i])
		pre++
for i in range(pre, n):
	if nums[i] == 1:
		swap(nums[pre], nums[i])
		pre++

但是如果要求只能遍历一遍，就只能用双指针遍历，大体思路是这样的：

  1. 双指针，p0指向0的区域，p1指向1的区域

  2. 因为我们是把0放在最左边，1放在中间，那就要注意，在交换后续的0时，很可能把之前排好的1换出去了

  3. 因此，如果遇到0并且p0 < p1（p0=p1说明还没有1出现，交换没有意义），不仅要交换nums[i], nums[p0],还要交换nums[i], nums[p1]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
        *// 0   0   1   1   2   0*

​        *//        p0      p1   i*

​        *// swap(nums[i], nums[p0])*

​        *// 0   0   0   1   2   1*

​        *//            p0  p1  i*

​        *// swap(nums[i], nums[1])*

​        *// 0   0   0   1   1  2*

​        *//            p0      p1*

代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Solution {
public:
    void sortColors(vector<int>& nums) {
        int p0 = 0, p1 = 0;
        for (int i = 0; i < nums.size(); i++) {
            if (nums[i] == 1) {
                swap(nums[p1], nums[i]);
                p1++;
                continue;
            }
            if (nums[i] == 0) {
                swap(nums[p0], nums[i]);
                if (p0 < p1) {
                    swap(nums[p1], nums[i]);
                }
                p0++, p1++;
            }
        }
    }       
};