題目資訊

• 題目編號: 75
• 題目連結: Sort Colors
• 主題: Array, Two Pointers, Sorting

題目敘述

給定一個包含 $n$ 個物件的陣列 nums，每個物件的顏色為紅、白或藍，請 就地 排序，使得相同顏色的物件相鄰，並且顏色順序為紅、白、藍。

我們將使用整數 0、1 和 2 來分別代表紅、白、藍色。

你必須在不使用庫的排序功能的情況下解決這個問題。

範例 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 \leq n \leq 300$
• $nums[i]$ 為 0、1 或 2。

進階: 你能想出一個僅用常數額外空間的一次遍歷演算法嗎？

直覺

這個問題是對具有固定數量不同元素的陣列進行排序的典型範例。根據限制條件，元素只能取三個值之一（0、1 或 2），我們可以利用計數的原則來解決這個排序問題。核心思想是計算每個不同元素的出現次數，然後根據這些計數重建陣列，以達到所需的排序順序。

解法

該解法採用了兩步驟的方法來高效地解決這個問題：

1. 計數出現次數：

   • 初始化一個大小為 3 的陣列 colorCount，將其值皆設為零，其中每個索引代表一種顏色（0, 1, 2）。
   • 遍歷給定的 nums 陣列，並在 colorCount 陣列中遞增每個遇到的元素的計數。
   • 在這一步後，colorCount[0]、colorCount[1] 和 colorCount[2] 將分別持有 nums 中紅色（0）、白色（1）和藍色（2）物件的數量。

2. 重建陣列：

   • 使用獲得的計數覆寫原始 nums 陣列：
     • 首先，將 nums 的前 colorCount[0] 個位置填入 0。
     • 接下來，填入之後的 colorCount[1] 個位置 1。
     • 最後，填入剩下的 colorCount[2] 個位置 2。
   • 這樣填充 nums 陣列，以使得所有的 0 都位於所有的 1 前面，然後是所有的 2，從而達到所需的排序順序。

此方法的時間複雜度為 $O(n)$，因為它遍歷陣列來計算元素數量並根據計數重填陣列需要單獨的遍歷。它使用 $O(1)$ 的額外空間來保存 colorCount 陣列，滿足了使用常量額外空間的限制。

程式碼

C++

class Solution {
public:
    void sortColors(vector<int>& nums) {
        // 陣列用來計算0、1和2的出現次數
        array<int, 3> colorCount = {0, 0, 0};

        // 計算輸入陣列中每個顏色的數量
        for (int color : nums) {
            colorCount[color]++;
        }

        // 使用已排序的顏色覆蓋原始陣列
        // 先用0填充陣列
        for (int i = 0; i < colorCount[0]; ++i) {
            nums[i] = 0;
        }

        // 然後用1填充
        for (int i = colorCount[0]; i < colorCount[0] + colorCount[1]; ++i) {
            nums[i] = 1;
        }

        // 最後用2填充
        for (int i = colorCount[0] + colorCount[1]; i < colorCount[0] + colorCount[1] + colorCount[2]; ++i) {
            nums[i] = 2;
        }
    }
};

Python

class Solution:
    def sortColors(self, nums):
        # 用於計數0、1和2出現次數的陣列
        colorCount = [0, 0, 0]

        # 計算輸入陣列中每個顏色的數量
        for color in nums:
            colorCount[color] += 1

        # 使用排序後的顏色覆蓋原始陣列
        # 首先用0填充陣列
        for i in range(colorCount[0]):
            nums[i] = 0

        # 接著用1填充
        for i in range(colorCount[0], colorCount[0] + colorCount[1]):
            nums[i] = 1

        # 最後用2填充
        for i in range(colorCount[0] + colorCount[1], colorCount[0] + colorCount[1] + colorCount[2]):
            nums[i] = 2

複雜度分析

• 時間複雜度: $O(n)$
  • 該演算法涉及兩個主要步驟，每個步驟都需要遍歷整個輸入陣列 nums。
  • 在第一個步驟中，它遍歷 nums 以計算每種顏色的出現次數。由於這個步驟對陣列中的每個元素進行一次處理，因此時間複雜度為 $O(n)$。
  • 在第二個步驟中，演算法使用在第一步中收集的計數覆蓋原始陣列為已排序的顏色。這步驟也需要遍歷 nums 並將 0s、1s 和 2s 的計數寫回陣列，因此同樣需要 $O(n)$ 的時間。
  • 由於這兩個步驟是依次進行的，總的時間複雜度為 $O(n) + O(n) = O(n)$。
• 空間複雜度: $O(1)$
  • 演算法使用一個固定大小為3的陣列 colorCount 來追蹤每種顏色（0，1，和2）的計數。colorCount 陣列的大小不依賴於輸入大小 n，因此空間複雜度為 $O(1)$，表示常數的空間使用。
  • 沒有使用和輸入大小成比例的額外空間；排序是通過直接修改輸入陣列 nums 原地完成的。