Pankaj Tanwar
Published on

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

––– views

Yay! It's the second day of my coding diary.

Problem of the day - Sort an array of 0s, 1s and 2s

Tag - Easy

So, the problem is, an array containing numbers 0,1, or 2 is given and I need to sort the array.

The catch here is, I can't use any in-build sorting algorithm here. I need to keep the time complexity o(n).

Initially, the problem seemed pretty complex but after some drilling, a simple approach came into my mind. As I only have 0,1 or 2 so why not store their counts and then just push values as per the count to new array.

Suppose, 0 came twice, 1 came thrice & 2 came once. So, I will just push values as per count. My final output would look like this - [0,0,1,1,1,2]

Code

#include <iostream>
#include<vector>
using namespace std;


int main()
{
    vector<int> list = {0,1,2,2,1,1};


    vector<int> count = {0,0,0};


    for(auto val: list) {
        count[val]++;
    }
    vector<int> res;


    int curr_index      = 0;
    int output_index    = 0;


    while(curr_index < count.size()) {
        if(count[curr_index] == 0) {
            curr_index++;
        } else {
            res.push_back(curr_index);
            count[curr_index] -= 1;
        }
    }
    for(auto val: list) {
        cout << val << " ";
    }
    return 0;
}

Time complexity - o(n)

Space Complexity - o(n)

Umm, It seems I am unnecessary keeping an extra array. Can I do it without using extra space (variables are not counted as extra space) ?

As I have calculated the count, I don't care what were the values in the original array so I can just update the same array, in place.

My modified code looks like this -

#include <iostream>
#include<vector>
using namespace std;


int main()
{
    vector<int> list = {0,1,2,2,1,1};


    vector<int> count = {0,0,0};


    for(auto val: list) {
        count[val]++;
    }


    int curr_index      = 0;
    int output_index    = 0;


    while(curr_index < count.size()) {
        if(count[curr_index] == 0) {
            curr_index++;
        } else {
            list[output_index] = curr_index;
            output_index++;
            count[curr_index] -= 1;
        }
    }
    for(auto val: list) {
        cout << val << " ";
    }
    return 0;
}

Time complexity - o(n)

Space Complexity - o(1)

You might like previous editions of my coding diary -

Tweet this article