题目描述

[EN | CN]

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

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

注意:

不能使用代码库中的排序函数来解决这道题。

示例:

1
2
输入: [2,0,2,1,1,0]
输出: [0,0,1,1,2,2]

进阶:

  • 一个直观的解决方案是使用计数排序的两趟扫描算法。首先,迭代计算出 0、1 和 2 元素的个数,然后按照 0、1、2 的排序,重写当前数组。
  • 你能想出一个仅使用常数空间的一趟扫描算法吗?

解法 1:2 次计数排序

实现一下题目中写的 2 次计数排序,统计所有数字的数量,然后重写该数组。

假设数组长度为 $n$,那么

  • 时间复杂度为 $O(n)$,遍历两遍;
  • 空间复杂度为 $O(1)$,常数级。

实现与结果如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        nums_elements = [0, 0, 0]
        for n in nums:
            nums_elements[n] += 1
        start = 0
        for n, elements in enumerate(nums_elements):
            for i in range(elements):
                nums[start + i] = n
            start += elements
  • 执行用时:40 ms,在所有 Python3 提交中击败了 62.25% 的用户。
  • 内存消耗:13.6 MB,在所有 Python3 提交中击败了 8.33% 的用户。

解法 2:三指针法

根据题目的提示,要求我们只能进行一次扫描,但是我们又需要进行三种颜色的重排,那么我们考虑用三个指针 ijk 分别指向 0、1、2 的尽头,其中:

  • i 指针从左往右,指向下一个将要与 j 交换的元素;交换后该元素变成 0,i 指针向右走。
  • j 指针从左往右,负责遍历,如果指向 1 则继续走;如果指向 0 或 2 则分别与 ij 交换。
  • k 指针从右往左,指向下一个将要与 j 交换的元素;交换后该元素变成 1,k 指针向左走。

一直遍历到 jk 指针相遇。

扫描的步骤如下:

  • jk 指针未相遇:
    • 如果当前 j 指针指向 0,那么 ij 指针指向的元素交换,都向右走一步;
    • 如果当前 j 指针指向 1,那么 j 指针向右走一步;
    • 如果当前 j 指针指向 2,那么 jk 指针指向的元素交换,k 向左走一步。

假设数组长度为 $n$,那么

  • 时间复杂度为 $O(n)$,遍历一遍;
  • 空间复杂度为 $O(1)$,常数级。

实现与结果如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        i, j, k = 0, 0, len(nums) - 1
        while j <= k:
            if nums[j] == 0:
                nums[i], nums[j] = 0, nums[i]
                i += 1
                j += 1
            elif nums[j] == 1:
                j += 1
            else:
                nums[j], nums[k] = nums[k], 2
                k -= 1
  • 执行用时:32 ms,在所有 Python3 提交中击败了 92.38% 的用户。
  • 内存消耗:13.7 MB,在所有 Python3 提交中击败了 8.33% 的用户。