Leetcode 75. Sort Colors

Leetcode 75. 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.

This problem is a sorting problem, actually. But the Quick Sort is not the best choice, since there are only 3 kinds of item in the array. Therefore, a special sorting algorithm could be carefully implemented in O(n) time.

Solution #1

Although in this approach, there are nested for/while loop, it is O(n) algorithm actually. The non_target_head index point goes through the whole array only in one pass. The key idea in this approach is placing the target number in the non_target_head location during each iteration. The non_target_head points to the first non-target number in current array. If a target number has been placed, then the non_target_head pointer moves forward to the next non-target item and never look back. The inner for-loop is to find the next target number to be placed in the location of non_target_head.

class Solution {
public:
    void sortColors(vector<int>& nums)
    {
        int non_target_head = 0;
        int size = nums.size();

        for(int target = 0; target < 3; target++)
        {
            // find the first non-target item's index
            while(non_target_head < size && nums[non_target_head] == target)
                non_target_head++;

            for(int i = non_target_head + 1; i < size; i++)
            {
                if(nums[i] == target)
                {
                    nums[i] = nums[non_target_head];
                    nums[non_target_head] = target;

                    // find next non-target item's index
                    while(non_target_head < size && nums[non_target_head] == target)
                        non_target_head++;

                    if(non_target_head == size)
                        break;
                }
            }
        }
    }
};

Solution #2

This is nicer solution. In this approach, we got three points red_index, white_index and blue_index. Each index represents the next target location to be placed (e.g., red_index is the location of next red target to be placed).

In the for-loop, there are 3 cases:

• when nums[i] == 2, then add one more 2 (e.g., blue) into array, and move blue_index 1 position forward.
• when nums[i] == 1, then add one more 1 (e.g., white) into array. Note: before we put the number 1 into array, we need to shift current number 2 one position left to make a gap.
• when nums[i] == 0, then add one more 0 (e.g., red) into array. And, of course, shift current number 2 and current number 1 one position left, respectively. Note: shift number 2 first, then shift number 1, because all the number 2 is after any number 1.

class Solution {
public:
    void sortColors(vector<int>& nums)
    {
        int red_index, white_index, blue_index;
        int size = nums.size();

        red_index = white_index = blue_index = 0;

        for(int i = 0; i < size; i++)
        {
            if(nums[i] == 0)
            {
                nums[blue_index++] = 2;
                nums[white_index++] = 1;
                nums[red_index++] = 0;
            }
            else if(nums[i] == 1)
            {
                nums[blue_index++] = 2;
                nums[white_index++] = 1;
            }
            else
            {
                nums[blue_index++] = 2;
            }
        }
    }
};