Pankaj Tanwar
Published on

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

––– views
Authors
•  Name
Pankaj Tanwar
@the2ndfloorguy

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 -