颜色分类 | 秋招知识库

题目链接 · 灵神原题解（署名来源）

不让用 $sort$ 吗？有意思……

技巧：O(1) 插入元素

假设现在有一个有序数组 $a = [0, 0, 1, 1, 2, 2]$ 。在 $a$ 中插入一个 $0$ ，同时保证 $a$ 是有序的，你会怎么做？

最暴力的想法是，把 $0$ 插在数组的最左边，原来的元素全体右移一位，得到 $[0, 0, 0, 1, 1, 2, 2]$ 。这样做是 $O (n)$ 的。

实际上，我们可以「狸猫换太子」：不是插入元素，而是修改元素！

对比一下插入前后：

插入前 $[0, 0, 1, 1, 2, 2]$ 。
插入后 $[0, 0, 0, 1, 1, 2, 2]$ 。

竖着看，其实只有三个位置上的数变了：

$a [2]$ 从 $1$ 改成 $0$ 。
$a [4]$ 从 $2$ 改成 $1$ 。
末尾新增一个 $2$ ，相当于 $a [6] = 2$ 。

怎么知道这些位置？

维护 $a$ 中 $0$ 的个数，即为改成 $0$ 的位置，记作 $p_{0}$ 。上例中 $p_{0} = 2$ 。把 $a [p_{0}]$ 改成 $0$ 。
维护 $a$ 中 $0$ 和 $1$ 的个数，即为改成 $1$ 的位置，记作 $p_{1}$ 。上例中 $p_{1} = 4$ 。把 $a [p_{1}]$ 改成 $1$ 。
末尾新增的位置记作 $i$ ，把 $a [i]$ 改成 $2$ 。

细节

如果 $a$ 中没有 $2$ 呢？上面第三步就错了。

比如 $a = [1]$ ，插入一个 $0$ ，结果为 $[0, 1]$ 。但如果按照上面三步走，第三步把 $a [1]$ 改成 $2$ ，得到的是错误的 $[0, 2]$ 。

要用很多 $if-else$ 特判这些特殊情况吗？

不需要，我们可以倒过来算：先把 $a [1]$ 改成 $2$ ，再把 $a [1]$ 改成 $1$ （覆盖），最后 $a [0]$ 改成 $0$ ，得到 $[0, 1]$ 。这种「覆盖」等价于「没有 $2$ 的时候不改成 $2$ 」。

如果插入的是 $1$ 呢？

跳过「把 $a [p_{0}]$ 改成 $0$ 」这一步。

如果插入的是 $2$ 呢？

只需要把 $a [i]$ 改成 $2$ 。

本题思路

对 $nums$ 执行插入排序，也就是对 $i = 0, 1, 2, \dots, n - 1$ 依次执行如下过程：

现在前缀 $nums [0]$ 到 $nums [i - 1]$ 是有序的，把 $nums [i]$ 插入这个有序前缀，从而把前缀 $nums [0]$ 到 $nums [i]$ 变成有序的。
算法执行完后， $nums$ 就是一个有序数组了。

Python 解法1C++ 解法1

python

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

cpp

// C++ 版待补充

C++ 解法2

cpp

class Solution {
public:
    void sortColors(vector<int>& nums) {
        int p0 = 0, p1 = 0;
        for (int i = 0; i < nums.size(); i++) {
            int x = nums[i];
            nums[i] = 2;
            if (x <= 1) {
                nums[p1++] = 1;
            }
            if (x == 0) {
                nums[p0++] = 0;
            }
        }
    }
};

复杂度分析

时间复杂度： $O (n)$ ，其中 $n$ 是 $nums$ 的长度。
空间复杂度： $O (1)$ 。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

技巧：O(1) 插入元素 ​

细节 ​

本题思路 ​

复杂度分析 ​

分类题单 ​

技巧：O(1) 插入元素

细节

本题思路

复杂度分析

分类题单