[Leetcode] Sort Colors

Sort Colors

 Given an array with n objects colored red, white or blue, sort them so that objects of the same color are adjacent, with the colors in the order red, white and blue. Here, we will use the integers 0, 1, and 2 to represent the color red, white, and blue respectively. 

Analysis: 
At first glance it feels so similar to 'Remove duplicates from a string'.  I wrote my own solution with 2 iterations.
Solution 1:
void sortColors(int A[], int n) {
     int i = 0, j = 0;   
    for(i = 0; i < n; i++)            -> 1st loop move all 0s to the front
    {
        if(A[i] == 0)
        {
            int temp = A[i];
            A[i] = A[j];
            A[j] = temp;
            j++;
        }
    }
    
    for(i = j; i < n; i++)           -> 2nd loop move all 1s to the front (after 0s)
    {
        if(A[i] == 1)
        {
            int temp = A[i];
            A[i] = A[j];
            A[j] = temp;
            j++;
        }
    }
}

For this specific problem (all elements fall into exactly 3 categories), it can be improved to single iteration, below is a sample from online source:

Solution 2:
void swap(int *a, int i, int j)
{
    int temp = a[i];
    a[i] = a[j];
    a[j] = temp;
}


void sortColors(int A[], int n) {

    
    if(n == 1)
      return;
      
    int i = 0, isOne = 0, isTwo = n-1;
    
    while(i < isTwo+1)
    {
        if(A[i] == 0)                     -> move 0s to the front
        {
            swap(A, i, isOne);
            isOne++;
            i++;
            continue;
        }
        if(A[i] == 2)                     -> move 2s to the end
        {
            swap(A, i, isTwo);
            isTwo--;                      -> after swapping, we don't move pointer isTwo cuz the new value can be 1,
            continue;                        we will compare it with i in next loop                          
            
        }
        else
        {
            i++;
        }
        
    }
      
}


Final thoughts:
For this specific problem solution 2 is better because of better time complexity. However if number of categories is k ( k > 3),  solution 1 is the only choice and time complexity would be O(kn).