SORT AN ARRAY OF 0'S ,1'S AND 2'S

Hey friends this is one of the easy problem but it is important for interview purpose:

problem: we have an array of 0's , 1's , 2's. we have to sort the array in o(n) time.

example: arr[]={2,1,0,2,0,2,1};
after sorting arr[] should be :{0,0,1,1,2,2,2};

Time Complexity:o(n)
Space Complexity;o(1)

Here is the code for the problem if u have any doubt let me know:

#include<bits/stdc++.h>
using namespace std;
int main(){
 int n,i,j;//size of array
 cout << "Enter size of array\n";
 cin >> n;

 int arr[n];
 cout << "Enter the elements of array:\n";
 for(i=0;i<n;i++){
  cin >> arr[i];
 }

 //I am taking 3 variable for the count of 0's and 1's and 2's.
 int count_0=0;
 int count_1=0;
 int count_2=0;

 for(i=0;i<n;i++){
  if(arr[i]==0){
   count_0++;
  }
  else if(arr[i]==1){
   count_1++;
  }
  else{
   count_2++;
  }
 }

 i=0;
 while(count_0>0){
  arr[i++]=0;
  count_0--;
 }

 while(count_1>0){
  arr[i++]=1;
  count_1--;
 }

 while(count_2>0){
  arr[i++]=2;
  count_2--;
 }

 cout << "finally the array is :\n";
 for(i=0;i<n;i++){
  cout << arr[i] << " ";
 }

 return 0;
}