Sort an array of zeros, ones and twos in linear time complexity

Algorithm

1. Set three variables low=0,mid=0, high=n-1 where n=length of input array
2. While (mid<=high), consider three cases on basis of the value of array [mid].
3. If array[mid] is 0
   Swap (array [low], array [mid]) since 0's should be at starting
   Increment low
   Increment mid
   Break statement
4. If array[mid] is 1
   Keep as it is since 1's should be at middle after being sorted
   Increment mid
   Break statement
5. If array[mid] is 2
   Swap (array [mid], array [high])
   Decrement high
   Break statement
6. Resulting array is sorted.

C++ program to sort an array of 0's, 1's and 2's in linear time complexity

#include<stdio.h>
#include<stdlib.h>

//function to swap by reference
void swap(int*a,int*b){
	int temp;
	temp=*b;
	*b=*a;
	*a=temp;
	return;
}
int* asort(int *a,int n){
	int low=0,mid=0,high=n-1;   //variables are set
	
	while(mid<=high){
		switch(a[mid]){
			case 0:     //if a[mid]==0
				//swap a[low] & a[mid], swapping by reference
				swap(&a[low],&a[mid]);  
				low++;      //increment low
				mid++;      //increment mid
				break;
			case 1:     //if a[mid]==1
				mid++;      //increment mid
				break;
			case 2:     //if a[mid]==2
				//swap a[mid] & a[high], swapping by reference
				swap(&a[mid],&a[high]);  
				high--;     //decrement high
				break;
		}
	}
	//returning adress of array(sorted)
	return a;   
}

int main() {
	int n;

	printf("enter no of array elements\n");
	//input array length
	scanf("%d",&n);

	int* a=(int*)malloc(sizeof(int)*n);

	printf("enter array elements\n");
	//input array elements
	for(int j=0;j<n;j++)      
		scanf("%d",&a[j]);
	//array is modified
	a=asort(a,n);   
	printf("after being sorted..............\n");
	//printing the sorted array
	for(int j=0;j<n-1;j++)   
		printf("%d ",a[j]);
	
	printf("%d\n",a[n-1]);

	return 0;
}

Output

First run:
enter no of array elements 
10 
enter array elements 
0
2
2
1
1
1
2
1
2
1
after being sorted.............. 
0 1 1 1 1 1 2 2 2 2 


Second run:
enter no of array elements 
10 
enter array elements 
2
2
2
2
0
0
0
0
1
1
after being sorted.............. 
0 0 0 0 1 1 2 2 2 2