Sort an array 0s, 1s and 2s (Dutch National Flag Problem)

Problem

Given an array containing only 0s, 1s and 2s. Write an efficient algorithm to sort the array.

Sample Input

[2, 0, 1, 1]

Sample Output

[0, 1, 1, 2]

Brute Force

You can directly use any sorting algorithm to do this task. This will take O(NLogN) time and at least O(N) auxiliary space. Can we find a better approach?

Efficient approach

This problem can be solved using a 3 pointer approach. Let’s see how! There are three pointers: start , idx and finish . start stores the first index that is not 0, finish stores the last index that is not 2 and idx will go from start to finish . If the element is 0, replace it with the element at index start and update start = start + 1 and idx = idx + 1 . If the element is 1, then idx = idx + 1 should be updated. If the element is 2, swap it with the element at index finish and update finish = finish – 1 . The time complexity of this approach is O(N) with O(1) auxiliary space.

C++ Programming

#include<bits/stdc++.h>
using namespace std;

void sortArray(vector<int>& nums) {
        int n=nums.size();
        if(n==1) return;
        int start=0;
        int finish=n-1;
    
        while(start<n && nums[start]==0){
            start++;
        }
        while(finish>=0 && nums[finish]==2){
            finish--;
        }
        int idx=start;
        while(idx<=finish){
            if(nums[idx]==0){
                swap(nums[start],nums[idx]);
                start++;
                idx++;
            }
            else if(nums[idx]==2){
                swap(nums[finish],nums[idx]);
                finish--;
            }
            else idx++;
        }
    }
int main(){
    vector<int> nums{1, 1, 0, 2};
    sortArray(nums);
    for(auto itr: nums) cout << itr<<" ";
}

C Programming

#include<stdio.h>
void sortArray(int nums[], int numsSize){
    int finish = numsSize -1;
    int start = 0;
    int idx = 0;
    while (idx <= finish)
    {
        if (nums[idx] == 0)
        {
            int save = nums[start];
            nums[start] = nums[idx];
            nums[idx] = save;
            idx++; start++;
        }
        else if (nums[idx] == 1)
        {
            idx++;
        }
        else
        {
            int temp = nums[finish];
            nums[finish] = nums[idx];
            nums[idx] = temp;
            finish--;
        }
    }
}
void main(){
    int nums[4] = {1, 1, 0, 2};
    sortArray(nums, 4);
    for(int i=0; i<4; i++) printf("%d ",nums[i]);
}

Python Programming

def sortArray(a):
        idx = 0
        start = 0
        finish = len(a) - 1
        while idx <= finish:
            if a[idx] == 0:
                a[idx], a[start] = a[start], a[idx]
                idx += 1
                start += 1
            elif a[idx] == 1:
                idx += 1
            else:
                a[idx], a[finish] = a[finish], a[idx]
                finish -= 1
        return a
        
l = [1, 1, 0, 2]
print(sortArray(l))

Output

[0, 1, 1, 2]

People are also reading:

• Merge Two Sorted Arrays in-place
• Subarray with Sum k
• Find Maximum Subarray Sum
• Longest subarray with sum as K
• Majority Element
• Find whether an array is subset of another array
• Rod Cutting Problem – Dynamic Programming
• Print a given matrix in spiral form
• Find Maximum Length Subarray Having an Equal Number of 0’s and 1’s
• Check for pairs in an array with a particular sum