75. 颜色分类

Posted 2021-09-17 炫云云

75. 颜色分类

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

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

示例 1：

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

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

输入：nums = [0]
输出：[0]

单指针

我们可以考虑对数组进行两次遍历。在第一次遍历中，我们将数组中所有的 0 交换到数组的头部。在第二次遍历中，我们将数组中所有的 1 交换到头部的 0 之后。此时，所有的 2 都出现在数组的尾部，这样我们就完成了排序。

具体地，我们使用一个指针 $\text{ptr}$ 表示「头部」的范围，$\text{ptr}$ 中存储了一个索引，$\text{ptr}$ 的初始值为 0 ，

在第一次遍历中，我们从左向右遍历整个数组，如果找到了 0，那么就需要将 0与 nums[ptr] 进行交换，并 ptr += 1。在遍历结束之后，所有的 0 都被交换到「头部」的范围，并且「头部」只包含 0 。

在第二次遍历中，我们从「头部」开始，从左向右遍历整个数组，如果找到了 1，那么就需要将 1 与 nums[ptr] 进行交换，并 ptr += 1。在遍历结束之后，所有的 1 都被交换到「头部」的范围，并且都在 0 之后，此时 22只出现在「头部」之外的位置，因此排序完成。

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        n = len(nums)
        ptr = 0
        for i in range(n):
            if nums[i] ==0:
                nums[i] ,nums[ptr] = nums[ptr] ,nums[i]
                ptr+=1
        for i in range(ptr,n):
            if nums[i] ==1:
                nums[i] ,nums[ptr] = nums[ptr] ,nums[i]
                ptr+=1

双指针

方法一需要进行两次遍历，那么我们是否可以仅使用一次遍历呢？我们可以额外使用一个指针，即使用两个指针分别用来交换 0 和 1。

具体地，我们用指针 $p_0$ 来交换 0 ，$p_1$ 来交换 $1$ ，初始值都为 0 。当我们从左向右遍历整个数组时：

如果找到了 1，那么将其与 $\text{nums}[p_1]$ 进行交换，并将 $p_1$ 向后移动一个位置，这与方法一是相同的；

如果找到了 0 ，那么将其与 $\text{nums}[p_0]$ 进行交换，并将 $p_0$ 向后移动一个位置。这样做是正确的吗？我们可以注意到，因为连续的 0 之后是连续的 1 ，因此如果我们将 0 与 $\text{nums}[p_0]$ 进行交换，那么我们可能会把一个 1 交换出去。

当 $p_0<p_1$ 时，我们已经将一些 1 连续地放在头部，此时一定会把一个 1 交换出去，导致答案错误。因此，如果 $p_0<p_1$ ，需要再将$\text{nums}[i]$ 与 $\text{nums}[p_1]$ 进行交换，其中 $i$是当前遍历到的位置，在进行了第一次交换后，$\text{nums}[i]$ 的值为 1 ，我们需要将这个 1 放到「头部」的末端。在最后，无论是否有 $p_0<p_1$ ，我们需要将 $p_0$ 和 $p_1$ 均向后移动一个位置，而不是仅将 $p_0$​​​ 向后移动一个位置。

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        n = len(nums)
        p0 = p1=0
        for i in range(n):
            if nums[i] ==1:
                nums[i] , nums[p1] = nums[p1] ,nums[i]
                p1 +=1
            elif nums[i] ==0:
                nums[i] , nums[p0] = nums[p0] ,nums[i]
                if p0 <p1:
                    nums[i] ,nums[p1] = nums[p1],nums[i]
                p0 += 1
                p1 += 1

双指针

与方法二类似，我们也可以考虑使用指针 $p_0$ 来交换 $0,p_2$ 来交换 2 。此时, $p_0$ 的初始值仍然为 0, 而 $p_2$ 的初始值为 $n-1_{\circ }$ 在遍历的过程中，我们需要找出所有的 0 交换至数组的头部，并且找出所
有的 2 交换至数组的尾部。

由于此时其中一个指针 $p_2$ 是从右向左移动的，因此当我们在从左向右遍历整个数组时，如果遍历到的位置超过了 $p_2$ ， 那么就可以直接停止遍历了。

具体地，我们从左向右遍历整个数组，设当前遍历到的位置为 $i$, 对应的元素为 $nums[i]$​​;

• 如果找到了 0 , 那么与前面两种方法类似，将其与 $\mathrm{nums}\left[p_0\right]$ 进行交换，并将 $p_0$ 向后移动一个位置;

• 如果找到了 2 ， 那么将其与 nums $\left[p_2\right]$ 进行交换，并将 $p_2$ 向前移动一个位置。

• 这样做是正确的吗? 可以发现，对于第二种情况，当我们将 nums $[i]$ 与 nums $\left[p_2\right]$ 进行交换之后，新的 nums $[i]$ 可能仍然是 2 ，也可能是 0 。然而此时我们已经结束了交换，开始遍历下一个元素 nums $[i+1]$, 不会再考虑 $nums[i]$ 了， 这样我们就会得到错误的答案。因此，当我们找到 2 时，我们需要不断地将其与 nums $\left[p_2\right]$ 进行交换，直到新的 $nums[i]$ 不为 2 。

  此 时，如果 nums $[i]$ 为 0 , 那么对应着第一种情况; 如果 $nums[i]$​​ 为 1, 那么就不需要进行任何后续的操作。

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        n = len(nums)
        p0, p2 = 0, n - 1
        i = 0
        while i <=p2:
            while i <=p2 and nums[i] ==2:
                nums[i], nums[p2] = nums[p2], nums[i]
                p2 -= 1
            if nums[i] == 0:
                nums[i], nums[p0] = nums[p0], nums[i]
                p0+=1
            i+=1

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        n = len(nums)
        p0, p2 = 0, n - 1
        if n < 2:
            return
        i = 0
        while i <=p2:
            if nums[i] == 0:
                nums[i], nums[p0] = nums[p0], nums[i]
                p0+=1
                i+=1
            elif nums[i] == 1:
                i += 1
            else:
                nums[i], nums[p2] = nums[p2], nums[i]
                p2 -= 1

参考

Krahets - 力扣（LeetCode） (leetcode-cn.com)