门|门模拟 2017 |第 48 题

本题是一个计算机科学中的经典问题，被称为“丹麦国旗问题”。题目描述如下：

在一个数组中有红、白、蓝三种颜色的球，设计一个算法将这三种球按照颜色顺序正确地排列。注意：这里仅是为了帮助理解题意，实际上该问题与当时的丹麦国旗无关。

具体要求如下：

• 算法应该尽可能地使用原地算法（不得使用额外的数组）；
• 算法应该具有线性时间复杂度。

解题思路

可以使用三个指针来实现。定义一个指向未分类部分开头的指针 $p$，一个指向红球部分末尾的指针 $r$，一个指向蓝球部分开头的指针 $b$。初始时，$p=0, r=0, b=n-1$（$n$ 为球的总数）。

然后遍历整个数组，根据节点上的颜色逐一处理：

• 如果是红色，将其交换到 $p$ 所指位置，然后 $p$ 和 $r$ 都向后一位；
• 如果是白色，将其保留在原地，然后 $p$ 向后一位；
• 如果是蓝色，将其交换到 $b$ 所指位置，然后 $b$ 向前一位。

如此往复直到 $p>b$。

C++代码实现

void sortColors(vector<int>& nums) {
    int n = nums.size();
    int p = 0, r = 0, b = n-1;
    while (p <= b) {
        if (nums[p] == 0) {
            swap(nums[p], nums[r]);
            p++;
            r++;
        } else if (nums[p] == 1) {
            p++;
        } else {
            swap(nums[p], nums[b]);
            b--;
        }
    }
}

Python 代码实现

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        n = len(nums)
        p, r, b = 0, 0, n-1
        while p <= b:
            if nums[p] == 0:
                nums[p], nums[r] = nums[r], nums[p]
                p += 1
                r += 1
            elif nums[p] == 1:
                p += 1
            else:
                nums[p], nums[b] = nums[b], nums[p]
                b -= 1