Problem Informations

• Problem Index: 75
• Problem Link: Sort Colors
• Topics: Array, Two Pointers, Sorting

Problem Description

Given an array nums with $n$ objects colored red, white, or blue, sort them in-place so that objects of the same color are adjacent, with the colors in the order red, white, and blue.

We will use the integers 0, 1, and 2 to represent the color red, white, and blue, respectively.

You must solve this problem without using the library’s sort function.

Example 1:

Input: nums = [2,0,2,1,1,0]
Output: [0,0,1,1,2,2]

Example 2:

Input: nums = [2,0,1]
Output: [0,1,2]

Constraints:

• $n == nums.length$
• $1 \leq n \leq 300$
• $nums[i]$ is either 0, 1, or 2.

Follow up: Could you come up with a one-pass algorithm using only constant extra space?

Intuition

The problem is a classic example of sorting an array with a fixed number of distinct elements. Given the constraints, where the elements can only take on one of three values (0, 1, or 2), we can approach this sorting problem by leveraging counting principles. The core idea is to count the occurrences of each distinct element and then reconstruct the array in the desired sorted order based on these counts.

Approach

The solution employs a two-step approach to solve the problem efficiently:

1. Counting Occurrences:

   • Initialize an array colorCount of size 3 to zero, where each index represents one of the colors (0, 1, 2).
   • Traverse the given nums array and increment the count of each encountered element in the colorCount array.
   • After this step, colorCount[0], colorCount[1], and colorCount[2] will respectively hold the number of red (0), white (1), and blue (2) objects in nums.

2. Reconstructing the Array:

   • Overwrite the original nums array using the counts obtained:
     • First, fill in the first colorCount[0] positions of nums with 0s.
     • Next, fill in the subsequent colorCount[1] positions with 1s.
     • Finally, fill in the remaining colorCount[2] positions with 2s.
   • This fills the nums array such that all 0s are followed by all 1s, which are then followed by all 2s, thus achieving the required sorted order.

This approach runs in $O(n)$ time complexity, as it makes a single pass through the array to count the elements and another pass to fill the array based on the counts. It uses $O(1)$ additional space for the colorCount array, satisfying the constraint of using constant extra space.

Code

C++

class Solution {
public:
    void sortColors(vector<int>& nums) {
        // Array to count occurrences of 0s, 1s, and 2s
        array<int, 3> colorCount = {0, 0, 0};

        // Count each color in the input array
        for (int color : nums) {
            colorCount[color]++;
        }

        // Overwrite the original array with sorted colors
        // First fill the array with 0s
        for (int i = 0; i < colorCount[0]; ++i) {
            nums[i] = 0;
        }

        // Then fill with 1s
        for (int i = colorCount[0]; i < colorCount[0] + colorCount[1]; ++i) {
            nums[i] = 1;
        }

        // Finally fill with 2s
        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):
        # Array to count occurrences of 0s, 1s, and 2s
        colorCount = [0, 0, 0]

        # Count each color in the input array
        for color in nums:
            colorCount[color] += 1

        # Overwrite the original array with sorted colors
        # First fill the array with 0s
        for i in range(colorCount[0]):
            nums[i] = 0

        # Then fill with 1s
        for i in range(colorCount[0], colorCount[0] + colorCount[1]):
            nums[i] = 1

        # Finally fill with 2s
        for i in range(colorCount[0] + colorCount[1], colorCount[0] + colorCount[1] + colorCount[2]):
            nums[i] = 2

Complexity

• Time complexity: $O(n)$
  • The algorithm involves two main steps that each iterate through the entire input array nums.
  • In the first step, it iterates over nums to count the occurrences of each color. This step takes $O(n)$ time as it processes each element of the array once.
  • In the second step, the algorithm overwrites the original array with sorted colors using the counts gathered in the first step. This step also involves iterating over nums and writing the counts of 0s, 1s, and 2s back into the array, which again takes $O(n)$ time.
  • Since both these steps are performed sequentially, the total time complexity is $O(n) + O(n) = O(n)$.
• Space complexity: $O(1)$
  • The algorithm uses a fixed-size array colorCount of size 3 to keep track of the count of each color (0, 1, and 2). The size of the colorCount array does not depend on the input size n, thus the space complexity is $O(1)$, indicating constant space usage.
  • No additional space proportional to the input size is used; the sorting is done in-place by modifying the input array nums directly.