Sort Colors - Leetcode Solution

Problem Link

💡 Step-by-Step Thought Process

Understand the problem: Sort an array of integers representing colors (0 for red, 1 for white, 2 for blue) in-place.
Initialize a counts array of size 3 to store the frequency of each color (0, 1, 2).
Iterate through each number (color) in the array nums and increment counts[color].
Extract the counts for red (R), white (W), and blue (B) from counts.
Modify nums in-place: set nums[0:R] to R zeros, nums[R:R+W] to W ones, and nums[R+W:] to B twos.

Code Solution

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        counts = [0, 0, 0]

        for color in nums:
            counts[color] += 1
        
        R, W, B = counts
        
        nums[:R]    = [0] * R
        nums[R:R+W] = [1] * W 
        nums[R+W:]  = [2] * B

#include <vector>

class Solution {
public:
    void sortColors(std::vector<int>& nums) {
        int counts[3] = {0, 0, 0};

        for (int color : nums) {
            counts[color]++;
        }

        int R = counts[0], W = counts[1], B = counts[2];

        for (int i = 0; i < R; i++) {
            nums[i] = 0;
        }
        for (int i = R; i < R + W; i++) {
            nums[i] = 1;
        }
        for (int i = R + W; i < nums.size(); i++) {
            nums[i] = 2;
        }
    }
};

public class Solution {
    public void sortColors(int[] nums) {
        int[] counts = new int[3];

        for (int color : nums) {
            counts[color]++;
        }

        int R = counts[0], W = counts[1], B = counts[2];

        for (int i = 0; i < R; i++) {
            nums[i] = 0;
        }
        for (int i = R; i < R + W; i++) {
            nums[i] = 1;
        }
        for (int i = R + W; i < nums.length; i++) {
            nums[i] = 2;
        }
    }
}

var sortColors = function(nums) {
    let counts = [0, 0, 0];

    for (let color of nums) {
        counts[color]++;
    }

    let R = counts[0], W = counts[1], B = counts[2];

    nums.fill(0, 0, R);
    nums.fill(1, R, R + W);
    nums.fill(2, R + W, R + W + B);
};

Detailed Explanation

Understanding the Sort Colors Problem

The "Sort Colors" problem is a classic interview question that involves sorting an array of integers where each element represents a color: 0 for red, 1 for white, and 2 for blue. The challenge is to sort the array in-place so that all 0s come first, followed by 1s, then 2s. This is also known as the Dutch National Flag problem, introduced by Edsger W. Dijkstra.

Conceptual Strategy

A brute-force approach might involve counting the frequency of each color and then writing those counts back to the array. While this works, it takes two passes over the data and isn’t as efficient as possible.

A more optimal solution uses the Dutch National Flag algorithm, which utilizes three pointers to sort the array in a single pass, in-place and with constant space.

Optimized Solution Using Dutch National Flag Algorithm

To solve this problem efficiently, we initialize three pointers: low, mid, and high. These will partition the array into three sections:

Indices before low will contain 0s.
Indices between low and mid will contain 1s.
Indices after high will contain 2s.

We iterate through the array with mid. If nums[mid] is 0, we swap it with nums[low] and move both pointers forward. If it’s 2, we swap it with nums[high] and move only high backward. If it’s 1, we simply move mid forward.

Time and Space Complexity

Time Complexity: O(n), where n is the length of the array. The algorithm only makes a single pass over the array.
Space Complexity: O(1), as the sorting is done in-place without using any additional data structures.

Why This Approach is Efficient

This method is preferred in interviews because it avoids the need for multiple passes or auxiliary storage. By using just a few pointers and conditional logic, the entire array can be sorted with minimal overhead, making it ideal for scenarios with memory or time constraints.

Conclusion

The Sort Colors problem not only tests your understanding of sorting and in-place algorithms but also demonstrates the power of pointer manipulation in solving classification problems. Mastering this problem can help you tackle a variety of in-place sorting challenges in real-world software development.

Get Personalized Support at AlgoMap Bootcamp 💡