【leetcode】75. 颜色分类（中等）（双指针、原地修改）

75. 颜色分类

问题描述

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

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

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

示例 1：

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

示例 2：

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

进阶：

• 你可以不使用代码库中的排序函数来解决这道题吗？
• 你能想出一个仅使用常数空间的一趟扫描算法吗？

思路&代码

这道题的要注意的就是不使用sort函数并且要进行原地排序。

使用双指针 $p_0$ 和 $p_1$ 分别记录当前应交换 0 和 1 的位置，即 $p_0$ 之前都是 0，$p_0$ 到 $p_1$ 之间都是 1，$p_1$ 及之后都是 2。注意，必有 $p_0\le p_1$。

遍历列表 nums，当遍历到第 i 个位置时：

• 当 $nums[i]=2$ 时，什么也不做。
• 当 $nums[i]=1$ 时，将 $nums[i]$ 与 $nums[p_1]$ 互换，并将 $p_1$ 后移一位。这时交换必然是值为 1 和 2 的交换，没有什么特殊情况。
• 当 $nums[i]=0$ 时，将 $nums[i]$ 与 $nums[p_0]$ 互换。需要注意的是，
  (1) 如果 $p_0=p_1$，则说明 还没遇到第二种情况（即$p_1$还未+1），则此时 i 后面的 1 还没换到前面来，也就是说 $nums[p_0]=nums[p_1]=2$，这时直接换完后把 $p_0$ 和 $p_1$ 后移一位即可，因为是值为 0 和 2 互换，这很正常；
  (2) 如果 $p_0<p_1$，则说明已经遇到过第二种情况（$p_1$已经后移，与$p_0$分开），此时 $nums[p_0]=1$，这时将 $nums[i]$ 与 $nums[p_0]$ 互换，则是值为 1 和 0 的互换，这时原本之前换到前面的 1 又换到后面去了，即换后$nums[i]=1$（先前是$nums[i]=0$），就需要再把这个 1 和前面的 2 换一下，也就是 $nums[i]$ 与 $nums[p_1]$ 互换，然后把 $p_0$ 和 $p_1$ 后移一位。

举个例子：

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

通过上面这个步骤每一步得到的结果为：

  这步以后其实就应该注意了，$p_0$ 和 $p_1$ 分开了，再遇到 0 的时候就不仅仅是交换 $nums[i]$ 与 $nums[p_0]$ 这么简单了。

  最后这步可以分成两个子步来解释，就是上面描述的第三种情况的(2)，$nums[i]$ 与 $nums[p_0]$交换后我们却把先前换到前面的1有放到后面去了，这里就要再把那个1放回到前面，于是就要再交换下$nums[i]$ 和 $nums[p_1]$：

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """ Do not return anything, modify nums in-place instead. """
        p0 = 0
        p1 = 0
        for i in range(len(nums)):
            if nums[i] == 0:
                nums[p0], nums[i] = nums[i], nums[p0]
                if p0 < p1:  
                    # 说明上步交换的是0和1，实际想要的是0和2交换，所以这里再把1(之前是0)和2交换
                    nums[p1], nums[i] = nums[i], nums[p1]
                p0 += 1
                p1 += 1
            elif nums[i] == 1:
                nums[p1], nums[i] = nums[i], nums[p1]
                p1 += 1
        return nums