Skip to main content

Two Pointer

Sort Colours Solution In C++/Java/Python/JS

Problem Description

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

We will use the integers 0, 1, and 2 to represent the colors:

0: Red
1: White
2: Blue
Sort Colours-Problem-Description

Requirements

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

Examples:

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

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


Constraints

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

This problem is a classic yet engaging challenge of in-place sorting. The goal is to rearrange an array of integers representing colors (red, white, blue) in the correct order efficiently. The key lies in partitioning the array into three groups using minimal swaps without relying on library functions.

Brute Force Approach

Intuition

The brute force approach provides a straightforward way to sort the array by categorizing and rearranging all the 0s, 1s, and 2s in order. The main idea is to count the number of occurrences of each color and then use this information to construct the sorted array.

First, we traverse the array and count how many times each color appears. For instance, if the array contains five 0s, three 1s, and two 2s, we record these counts in variables. This step gives us an overview of the distribution of colors in the array.

Next, we use the counts to overwrite the array. Starting from the beginning, we fill it with the correct number of 0s, followed by the correct number of 1s, and finally the correct number of 2s. Each color is placed consecutively, ensuring the array ends up sorted. 

By dividing the task into counting and placement, this method simplifies the problem into manageable steps, making it easy to implement and understand.

Approach

Step 1: Count the occurrences of 0s, 1s, and 2s

First, initialize three variables to store the count of 0s, 1s, and 2s. These variables will track how many times each color (0, 1, and 2) appears in the array. Start by iterating through the array and update the counts accordingly.

Step 2: Overwrite the array with sorted values

After counting the occurrences, the next step is to overwrite the original array. Starting from the beginning of the array, place the correct number of 0s, followed by the 1s, and finally the 2s. The sorted array should have all the 0s first, followed by 1s, and then 2s.

Step 3: Return the sorted array

After the array has been properly overwritten, it is now sorted. The array is modified in-place, so there is no need to return anything. However, you can return the nums array if required.

Let us understand this with a video ,

0:00
/1:20

Sort Colours-Brute-Force-Approach

Sort Colours Dry run - BruteForce

Input:nums = [2, 0, 2, 1, 1, 0]
Step 1: Count the occurrences of 0s, 1s, and 2s
Initialize variables count0 = 0, count1 = 0, and count2 = 0.
Iterate through the array:

  • nums[0] = 2: Increment count2 → count2 = 1
  • nums[1] = 0: Increment count0 → count0 = 1
  • nums[2] = 2: Increment count2 → count2 = 2
  • nums[3] = 1: Increment count1 → count1 = 1
  • nums[4] = 1: Increment count1 → count1 = 2
  • nums[5] = 0: Increment count0 → count0 = 2

Final counts: count0 = 2, count1 = 2, count2 = 2.

Step 2: Overwrite the array with sorted values
Now overwrite nums starting with the counts:

  • Place count0 (2 zeros): nums = [0, 0, 2, 1, 1, 0]
  • Place count1 (2 ones): nums = [0, 0, 1, 1, 1, 0]
  • Place count2 (2 twos): nums = [0, 0, 1, 1, 2, 2]

Step 3: Return the sorted array

  • The array is now sorted: nums = [0, 0, 1, 1, 2, 2].

This process modifies the array in-place and the result is the sorted array [0, 0, 1, 1, 2, 2].

Sort Colours Code for All Languages - BruteForce
C++
// Function to sort colors using counting approach
void sortColors(vector<int>& nums) {
    // Initialize counters for 0s, 1s, and 2s
    int count0 = 0, count1 = 0, count2 = 0;

    // Step 1: Count occurrences of 0s, 1s, and 2s
    for (int num : nums) {
        // If the number is 0, increment count0
        if (num == 0) count0++;
        // If the number is 1, increment count1
        else if (num == 1) count1++;
        // If the number is 2, increment count2
        else count2++;
    }

    // Step 2: Overwrite array with sorted values
    int index = 0;
    
    // Fill the array with 0s based on count0
    while (count0--) nums[index++] = 0;
    
    // Fill the array with 1s based on count1
    while (count1--) nums[index++] = 1;
    
    // Fill the array with 2s based on count2
    while (count2--) nums[index++] = 2;
}

Sort Colours Code in Java - BruteForce
public class Solution {
    // Function to sort colors using counting approach
    public void sortColors(int[] nums) {
        // Initialize counters for 0s, 1s, and 2s
        int count0 = 0, count1 = 0, count2 = 0;

        // Step 1: Count occurrences of 0s, 1s, and 2s
        for (int num : nums) {
            // If the number is 0, increment count0
            if (num == 0) count0++;
            // If the number is 1, increment count1
            else if (num == 1) count1++;
            // If the number is 2, increment count2
            else count2++;
        }

        // Step 2: Overwrite array with sorted values
        int index = 0;
        
        // Fill the array with 0s based on count0
        while (count0-- > 0) nums[index++] = 0;
        
        // Fill the array with 1s based on count1
        while (count1-- > 0) nums[index++] = 1;
        
        // Fill the array with 2s based on count2
        while (count2-- > 0) nums[index++] = 2;
    }

Sort Colours Code in Python - BruteForce
# Function to sort colors using counting approach
def sortColors(nums):
    # Initialize counters for 0s, 1s, and 2s
    count0 = count1 = count2 = 0

    # Step 1: Count occurrences of 0s, 1s, and 2s
    for num in nums:
        # If the number is 0, increment count0
        if num == 0:
            count0 += 1
        
        # If the number is 1, increment count1
        elif num == 1:
            count1 += 1
        
        # If the number is 2, increment count2
        else:
            count2 += 1

    # Step 2: Overwrite array with sorted values
    index = 0
    
    # Fill the array with 0s based on count0
    for _ in range(count0):
        nums[index] = 0  # Fill 0s
        index += 1
    
    # Fill the array with 1s based on count1
    for _ in range(count1):
        nums[index] = 1  # Fill 1s
        index += 1
    
    # Fill the array with 2s based on count2
    for _ in range(count2):
        nums[index] = 2  # Fill 2s
        index += 1

Sort Colours Code in JavaScript - BruteForce
// Function to sort colors using counting approach
function sortColors(nums) {
    let count0 = 0, count1 = 0, count2 = 0;

    // Step 1: Count occurrences of 0s, 1s, and 2s
    for (let num of nums) {
        if (num === 0) count0++;
        else if (num === 1) count1++;
        else count2++;
    }

    // Step 2: Overwrite array with sorted values
    let index = 0;
    while (count0--) nums[index++] = 0;  // Fill 0s
    while (count1--) nums[index++] = 1;  // Fill 1s
    while (count2--) nums[index++] = 2;  // Fill 2s
}

// User input: Size of the array
let n = parseInt(prompt()); 

// Create an array to hold the values
let nums = [];
for (let i = 0; i < n; i++) {
    nums.push(parseInt(prompt()));
}

// Call the function to sort the array
sortColors(nums);

// Print the sorted array
console.log(nums);

Time Complexity : O(N)

Counting occurrences:The algorithm iterates through the array once to count the occurrences of 0s, 1s, and 2s. This takes O(n), where nn is the length of the array.

Overwriting the array:The array is then overwritten based on the counts, which also requires iterating through the array once. This takes another O(n).

Thus, the total time complexity is O(n)+O(n)=O(n) .

Space Complexity : O(1)

Auxiliary space:The algorithm uses a constant amount of extra space for the count variables (count0, count1, count2), which is O(1).

Total space:The array nums is modified in-place, so no extra space is used for storing another array. Hence, the total space complexity is O(1).

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

Yes, we can sort the array in one pass using three pointers. The first pointer places the 0s, the second pointer keeps track of the current element, and the third pointer places the 2s. This method sorts the array efficiently in a single pass without using extra space.

Approach : Three Pointers (One-pass algorithm)

Intuition

Imagine you have a box full of balls in three colors: red, white, and blue, and your goal is to neatly arrange them so all the red balls come first, followed by the white ones, and finally the blue ones. At first, the colors are all jumbled up, but you can organize them quickly using a clever method. Instead of picking each ball and thinking about where it should go, you can simplify the process by focusing on just three tasks.

First, picture dividing the box into three sections: a spot for red balls at the beginning, a place for blue balls at the end, and a middle area where the white balls will naturally settle. To do this efficiently, you just need three imaginary markers. One marker will show where the next red ball should go, another will show where the next blue ball should go, and the third marker will help you examine each ball as you go through the box.

As you look at each ball, it’s easy to decide what to do. If it’s red, you swap it with the ball currently in the red section, expanding that area. If it’s blue, you move it to the blue section at the end of the box by swapping it with the ball in that spot, shrinking the blue section. If it’s white, you simply move past it, as white balls belong right in the middle where they already are. With each step, the box starts looking more organized, and you’re making progress without needing to backtrack.

By the time you’ve examined every ball, the job is done! All the reds are neatly at the start, the blues are at the end, and the whites fill the middle without any extra effort.

In the same way, we can think about sorting an array that contains only three distinct values—let’s say 0, 1, and 2. Imagine these numbers as representing the colors: 0 as white, 1 as red, and 2 as blue. The goal is to rearrange the array so that all the 0s come first, followed by the 1s, and then the 2s.

Approach

Step 1: Initialize Three Pointers
  • Set start to 0 (tracks the position for the next 0).
  • Set mid to 0 (used to traverse the array).
  • Set end to nums.length - 1 (tracks the position for the next 2).
Step 2: Traverse the Array
  • Use a while loop with the condition mid <= end to traverse the array.
Step 3: Check the Value at nums[mid]
  • Examine the value at nums[mid] to decide the operation.
Step 4: Handle nums[mid] == 0
  • Swap nums[mid] with nums[start].
  • Increment both start and mid pointers to place the 0 at its correct position.
Step 5: Handle nums[mid] == 1
  • Simply increment the mid pointer, as 1s naturally fall in the middle.
Step 6: Handle nums[mid] == 2
  • Swap nums[mid] with nums[end].
  • Decrement the end pointer to place the 2 at its correct position.
  • Do not increment mid immediately, as the swapped value needs to be checked again.
Step 7: Repeat Until Sorted

Continue the traversal until mid > end. At this point, the array is sorted with all 0s at the beginning, 1s in the middle, and 2s at the end.

Let us understand this with a video ,

0:00
/1:01

Sort Colours-Optimal-Approach

3Colours Dry Run - Three Pointer

Input: nums = [2, 0, 2, 1, 1, 0]
Initial Setup:

  • start = 0, mid = 0, end = 5
  • Initial array: [2, 0, 2, 1, 1, 0]

Step-by-Step Process:
Iteration 1 (mid = 0):

  • nums[mid] = 2
  • Swap nums[mid] and nums[end].
  • Updated array: [0, 0, 2, 1, 1, 2]
  • Pointers: start = 0, mid = 0, end = 4

Iteration 2 (mid = 0):

  • nums[mid] = 0
  • Swap nums[mid] and nums[start].
  • Increment both start and mid.
  • Updated array: [0, 0, 2, 1, 1, 2]
  • Pointers: start = 1, mid = 1, end = 4

Iteration 3 (mid = 1):

  • nums[mid] = 0
  • Swap nums[mid] and nums[start].
  • Increment both start and mid.
  • Updated array: [0, 0, 2, 1, 1, 2]
  • Pointers: start = 2, mid = 2, end = 4

Iteration 4 (mid = 2):

  • nums[mid] = 2
  • Swap nums[mid] and nums[end].
  • Decrement end.
  • Updated array: [0, 0, 1, 1, 2, 2]
  • Pointers: start = 2, mid = 2, end = 3

Iteration 5 (mid = 2):

  • nums[mid] = 1
  • Increment mid.
  • Updated array: [0, 0, 1, 1, 2, 2]
  • Pointers: start = 2, mid = 3, end = 3

Iteration 6 (mid = 3):

  • nums[mid] = 1
  • Increment mid.
  • Updated array: [0, 0, 1, 1, 2, 2]
  • Pointers: start = 2, mid = 4, end = 3

Final Output:

  • Array after sorting: [0, 0, 1, 1, 2, 2].
Sort Colours Code for All Languages - Three Pointer
C++
// Function to sort colors using three-pointer approach
void sortColors(vector<int>& nums) {
    // Initialize three pointers: start, mid, and end
    int start = 0, mid = 0, end = nums.size() - 1;

    // Step 2: Traverse the array
    while (mid <= end) {
        // If the element is 0, swap it to the start position
        if (nums[mid] == 0) {
            // Step 4: Swap and move pointers for 0s
            swap(nums[start], nums[mid]);
            start++;
            mid++;
        } 
        // If the element is 1, just move the mid pointer
        else if (nums[mid] == 1) {
            // Step 5: Move mid pointer for 1s
            mid++;
        } 
        // If the element is 2, swap it to the end position
        else {
            // Step 6: Swap and move end pointer for 2s
            swap(nums[mid], nums[end]);
            end--;
        }
    }
}

Sort Colours code in Java - Three Pointer
public class SortColors {

    // Function to sort colors using three-pointer approach
    public static void sortColors(int[] nums) {

        // Initialize three pointers: start, mid, and end
        int start = 0, mid = 0, end = nums.length - 1;

        // Step 2: Traverse the array
        while (mid <= end) {

            // If the element is 0, swap it to the start position
            if (nums[mid] == 0) {

                // Step 4: Swap and move pointers for 0s
                int temp = nums[start];
                nums[start] = nums[mid];
                nums[mid] = temp;
                start++;
                mid++;
            }
            // If the element is 1, just move the mid pointer
            else if (nums[mid] == 1) {

                // Step 5: Move mid pointer for 1s
                mid++;
            }
            // If the element is 2, swap it to the end position
            else {

                // Step 6: Swap and move end pointer for 2s
                int temp = nums[mid];
                nums[mid] = nums[end];
                nums[end] = temp;
                end--;
            }
        }
    }

Sort Colours code in Python - Three Pointer
# Function to sort colors using three-pointer approach
def sortColors(nums):
    start, mid, end = 0, 0, len(nums) - 1

    # Step 2: Traverse the array
    while mid <= end:
        if nums[mid] == 0:
            
            # Step 4: Swap and move pointers for 0s
            nums[start], nums[mid] = nums[mid], nums[start]
            start += 1
            mid += 1
        elif nums[mid] == 1:
            
            # Step 5: Move mid pointer for 1s
            mid += 1
        else:
            
            # Step 6: Swap and move end pointer for 2s
            nums[mid], nums[end] = nums[end], nums[mid]
            end -= 1

Sort Colours code in JavaScript - Three Pointer
// Function to sort colors using three-pointer approach
function sortColors(nums) {
    let start = 0, mid = 0, end = nums.length - 1;

    // Step 2: Traverse the array
    while (mid <= end) {
        if (nums[mid] === 0) {
            
            // Step 4: Swap and move pointers for 0s
            [nums[start], nums[mid]] = [nums[mid], nums[start]];
            start++;
            mid++;
        } else if (nums[mid] === 1) {
            
            // Step 5: Move mid pointer for 1s
            mid++;
        } else {
            
            // Step 6: Swap and move end pointer for 2s
            [nums[mid], nums[end]] = [nums[end], nums[mid]];
            end--;
        }
    }
}

Time Complexity: O(N)

The algorithm makes a single pass through the array, where each element is processed exactly once. The operations performed during each step (comparison and swapping) are constant time operations.

  • Number of operations: O(n), where nn is the length of the array.

Thus, the time complexity is O(n).

Space Complexity: O(1)

Auxiliary Space:The algorithm uses three pointers: start, mid, and end. These are constant space.
Auxiliary space complexity: O(1).

Total Space:Since no extra space beyond the input array and a few pointers is used, the total space complexity is also O(N).

Learning Tip

Now we have successfully tackled this problem, let's try these similar problems.

Given the head of a linked list, return the list after sorting it in ascending order.

Given an integer array nums, reorder it such that nums[0] < nums[1] > nums[2] < nums[3]....
You may assume the input array always has a valid answer.