Pankaj Tanwar
Published on

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

Authors

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 -