題目網址：https://leetcode.cn/problems/sort-colors/

題意：給一包含紅色、白色和藍色, 總共 n 個元素的 array nums, 原地對 nums 進行排序, 使得相同顏色的元素相鄰, 並按照紅、白、藍的順序排列。

分別使用 0、1、2 代表紅、白、藍。

注意：不能使用 library 中的 sort 函式。

進階：設計 $O(1)$ space 的一次掃描演算法。

Solution 1：

想法：第一次掃描先計算 0、1、2 的個數, 第二次掃描則按照個數填充 nums

class Solution {
public:
    void sortColors(vector<int>& nums) {
        vector<int> colors(3, 0);
        for (const auto& num : nums) {
            ++colors[num];
        }

        int cur = 0;
        for (int i = 0; i < 3; ++i) {
            for (int j = 0; j < colors[i]; ++j) {
                nums[cur++] = i;
            }
        }
    }
};
  • time：$O(n)$ ➔ $O(2n)$, 因為有 2 次掃描
  • space：$O(1)$ ➔ 撇除要返回的 res, 只需常數空間

Solution 2：

想法：利用 Two Pointers + Partition, 用 ptr i 來遍歷 nums, 定義並維護以下三個區間：

  • [0, left - 1] 皆為 0：所有在 left 以左的元素皆為 0
  • [left, i - 1] 皆為 1：left ~ i - 1 之間的元素皆為 1
  • [right + 1, n - 1] 皆為 2：所有在 right 以右的元素皆為 2

特殊情況：當 nums[i] == 2, nums[i] 要和 nums[right] 做 swap(), 但我們不知道 swap() 前原本 nums[right] 的元素是什麼, 因此 swap() 後不能直接將 i 往下一位移動, 而是應該在下一輪繼續判斷 nums[i]

  • swap() 前：

  • swap() 後：假設原先 nums[right] 為 0, 則 swap 後 nums[i] = 0, 若此時將 i 往下一位移動, 則那個 swap 的 0 不會被擺在正確的位置上

為什麼 nums[i] == 0 不是特殊情況：因為我們維護區間 [left, i - 1] 皆為 1, 也就是說 nums[left] 會指著 1, 和 nums[left] 做 swap() 後 nums[i] 變成 1, i + 1 後在下一輪中仍滿足 [left, i - 1], 因此可以直接往後一位

class Solution {
public:
    void sortColors(vector<int>& nums) {
        const int n = nums.size();
        int left = 0, right = n - 1;

        // 因為維護 right 以右皆為 2, 所以當 i > right 時結束
        for (int i = 0; i <= right; ++i) {
            if (nums[i] == 0) {
                swap(nums[i], nums[left++]);
            } else if (nums[i] == 2) {
                swap(nums[i--], nums[right--]); // i-- : 不讓 i 往下一位移動
            }
        }
    }
};
  • time：$O(n)$ ➔ for loop
  • space：$O(1)$ ➔ 只需常數空間