Sort an array of 0's, 1's and 2's

Hero Image

DT

Dhaval Trivedi

Co-founder, Airtribe

Problem Overview

The problem "Sort Colors" is a classic algorithm problem that requires sorting an array of integers representing colors, using an in-place algorithm without utilizing any library sort functions. The integers in the array represent the colors: 0 for red, 1 for white, and 2 for blue. The objective is to sort the array so that all 0s are grouped first, followed by 1s, and all 2s at the end.

Understanding the Problem

In this problem, we are given an array of n integers where each integer is either 0, 1, or 2. We need to sort the array such that integers of the same kind are adjacent, using a single-pass algorithm with constant space limitations. Conceptually, this is similar to the "Dutch National Flag" problem proposed by Edsger Dijkstra.

Example

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

Key DSA Concepts and Theory

Arrays

Arrays are a fundamental data structure that are commonly used when elements need to be accessed by a numerical index. In this problem, the array data structure is the main subject of the operation.

Sorting

Sorting is the process of arranging data in a particular order. In this problem, we're required to sort elements in increasing order, but we must do so in a single pass with constant space, which prevents using typical sorting algorithms like QuickSort or MergeSort.

Dutch National Flag Problem

The "Dutch National Flag" problem specifically deals with arranging elements of three distinct types in a single pass. It is highly applicable to the given problem, leveraging an algorithm that involves partitioning the array into three sections during iteration—one each for 0, 1, and 2.

Solution Approach

The optimal approach to solve this problem leverages the Dutch National Flag algorithm, which efficiently sorts the array in a single traversal with O(n) time complexity and O(1) space complexity.

Step-by-Step Explanation

  1. Initialization:

    • Three pointers are used: low, mid, and high.
    • low and mid start at the beginning of the array, while high starts at the end.
  2. Iteration:

    • Traverse through the array using the mid pointer.
    • When nums[mid] is 0: Swap nums[mid] with nums[low], increment low and mid.
    • When nums[mid] is 1: Simply increment mid.
    • When nums[mid] is 2: Swap nums[mid] with nums[high], decrement high.
  3. Termination:

    • The iteration stops when mid exceeds high, ensuring all elements are sorted.

Pseudocode

low = 0
mid = 0
high = len(nums) - 1

while mid <= high:
    if nums[mid] == 0:
        swap(nums[low], nums[mid])
        low++
        mid++
    elif nums[mid] == 1:
        mid++
    else:
        swap(nums[mid], nums[high])
        high--

Code: C++

#include <vector>
using namespace std;

void sortColors(vector<int>& nums) {
    int low = 0, mid = 0, high = nums.size() - 1;
    
    while(mid <= high) {
        if(nums[mid] == 0) {
            swap(nums[low++], nums[mid++]);
        } else if(nums[mid] == 1) {
            mid++;
        } else {
            swap(nums[mid], nums[high--]);
        }
    }
}

Code: Java

class Solution {
    public void sortColors(int[] nums) {
        int low = 0, mid = 0, high = nums.length - 1;
        
        while(mid <= high) {
            if(nums[mid] == 0) {
                swap(nums, low++, mid++);
            } else if(nums[mid] == 1) {
                mid++;
            } else {
                swap(nums, mid, high--);
            }
        }
    }
    
    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}

Time and Space Complexity Analysis

  • Time Complexity: O(n), where n is the number of elements in the array. Each element is inspected exactly once.
  • Space Complexity: O(1), as the algorithm uses a constant amount of extra space.

Common Mistakes to Avoid

  1. Incorrect Pointer Management: Ensure that mid is always incremented when required, particularly after handling an element.
  2. Infinite Loops: Properly manage pointers and conditions to avoid loops running indefinitely.

Similar Problems on LeetCode

Additional Resources and References

By understanding and applying the Dutch National Flag algorithm, this problem can be solved efficiently within the constraints given, enhancing both problem-solving and algorithmic understanding.