【数组&双指针】LeetCode 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]
 

提示：

n == nums.length
1 <= n <= 300
nums[i] 为 0、1 或 2
 

进阶：

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

【分析】

本题是经典的荷兰国旗🇳🇱问题，根据题目提示，我们可以统计出数组中0，1，2的个数，再根据它们的数量，重写整个数组。这种方法较为简单，也很容易想到。而这里会介绍两种基于指针进行交换的方法。

方法一：单指针 

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

具体地，我们使用一个指针ptr表示头部的范围，ptr中存储了一个整数，表示数组nums从位置0到位置ptr-1都属于头部。ptr的初始值为0，表示还没有数处于头部。

在第一次遍历中，我们从左向右遍历整个数组，如果找到了0，那么就需要将0与头部位置的元素进行交换，并将头部向后扩充一个位置。在遍历结束之后，所有的0都被交换到头部的范围，并且头部只包含0。

在第二次遍历中，我们从头部开始，从左向右遍历整个数组，如果找到了1，那么就需要将1与头部位置的元素进行交换，并将头部向后扩充一个位置，在遍历结束后，所有的1都被交换到头部的范围，并且都在0之后，此时所有的2都分布在头部之外的位置。排序完成。

时间复杂度：O(n)，n为数组nums的长度。

空间复杂度：O(1)。

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        n = len(nums)
        ptr = 0 # 定义指针ptr表示头部的范围，ptr中用于存储一个数组，表示数组nums从位置0到位置ptr-1都属于头部。其初始值为0，表示还没有数字处于头部之中。
        # 将数字0置于头部。
        for i in range(n):
            if nums[i] == 0:
                nums[i], nums[ptr] = nums[ptr], nums[i]
                ptr += 1
        # 将数字1置于头部，这个头部范围在0之后。
        for i in range(ptr, n):
            if nums[i] == 1:
                nums[i], nums[ptr] = nums[ptr], nums[i]
                ptr += 1

方法二 双指针

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

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

（1）若找到了1，那么将其与nums[p1]进行交换，并将p1向后移动一位，这与方法一相同。

（2）如果找到了0，那么将其与nums[p0]进行交换，并将p0向后移动一位，这样做是正确的吗？我们可以注意到，因为连续的0之后是连续的1，因此如果我们将0与nums[p0]进行交换，那么我们可能会把一个1交换出去。当p0<p1时，我们已经将一些1连续地放在头部，此时一定会把一个1交换出去，导致答案错误。因此，如果p0<p1，那么我们需要再将nums[i]与nums[p1]交换，其中i是当前遍历到的位置，在进行了第一次交换后，nums[i]的值为1，我们需要将这个1放到头部的末端。在最后，无论是否有p0<p1，我们都需要将p0和p1均向后移动一个位置，而不是只将p0向后移动一个位置。

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        n = len(nums)
        p0, p1 = 0, 0
        # 将数字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]
                # 当p0 < p1，
                if p0 < p1:
                    nums[i], nums[p1] = nums[p1], nums[i]
                p0 += 1
                p1 += 1

# 复杂度分析
# 时间复杂度：O(n)，其中 n 是数组nums 的长度。
# 空间复杂度：O(1)。

方法三：双指针

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

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

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

（1）若找到了0，那么与前面两种做法类似，将其与nums[p0]进行交换，并将p0向后移动一位；

（2）若找到了2，那么将其与nums[p2]进行交换，并将p2向前移动一个位置。

这样做对吗？可以发现，对于第二种情况，当我们将nums[i]与nums[p2]进行交换后，新的nums[i]可能仍然是2，也可能是0。然而此时我们已经结束了交换，开始遍历下一个元素nums[i+1]，不会再考虑nums[i]了，这样我们就会得到错误的答案。

因此当我们找到2时，需要不断地将其与nums[p2]进行交换，直到新的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