[Leetcode 97] 75 Sort Colors

Problem:

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.

Note:
You are not suppose to use the library's sort function for this problem.

 

Follow up:
A rather straight forward solution is a two-pass algorithm using counting sort.
First, iterate the array counting number of 0's, 1's, and 2's, then overwrite array with total number of 0's, then 1's and followed by 2's.

Could you come up with an one-pass algorithm using only constant space?

 

Analysis:

This problem can be solved via scanning and counting. After the first scan, we can know the number of each colored objects, then we just need to put corresponding number of colored objects into the array. The complexity is O(2n).

But the follow up requires O(n) time complexity and constant space algorithm. So we can use in-place swapping to solve the problem. Use idx i points to the 0's next to place position and idx j points to 2's next to place position. Then scan from 0 to j. 

If A[cur] == 0, swap A[cur] and A[i]. And increase i and cur;

If A[cur] == 2, swap A[cur] and A[j]. And decrease j;

If A[cur] == 1, do nothing. And increase cur.

Only one-pass is needed in this method.

 

Code:

Two-pass counting method

class Solution {
public:
    void sortColors(int A[], int n) {
        // Start typing your C/C++ solution below
        // DO NOT write int main() function
        int r = 0, w = 0, b = 0;
        
        for (int i=0; i<n; i++) {
            if (A[i] == 0)
                r++;
            else if (A[i] == 1)
                w++;
            else
                b++;
        }
        
        for (int i=0; i<n; i++) {
            if (i<r)
                A[i] = 0;
            else if (i>=r && i<r+w)
                A[i] = 1;
            else
                A[i] = 2;
        }
        
    }
};

 

One-pass swap method

class Solution {
public:
    void sortColors(int A[], int n) {
        // Start typing your C/C++ solution below
        // DO NOT write int main() function
        int low = 0, high = n-1;
        
        for (int i=0; i<=high; ) {
            if (A[i] == 0) {
                swap(A[i], A[low]);
                low++;
                i++;
            }
            else if (A[i] == 2) {
                swap(A[i], A[high]);
                high--;
            }
            else {
                i++;
            }
        }
        
    }
    
    void swap(int &a, int &b) {
        int tmp = a;
        a = b;
        b = tmp;
    }
};