0075. Sort Colors

Medium | Array + Two Pointer | 16 ms (93.28%), 13.2 MB (99.39%)

❓ Problem Statement

Source: LeetCode - Sort Colors GitHub: Solution / Performance

Given an array nums with n objects colored red, white, or blue, sort them in-place so that objects of the same color are adjacent, with the colors in the order red, white, and blue.

We will use the integers 0, 1, and 2 to represent the color red, white, and blue, respectively.

You must solve this problem without using the library's sort function.

Follow up: Could you come up with a one-pass algorithm using only constant extra space?

✍🏻 Constraints & Example

Constraints:

• n == nums.length

• 1 <= n <= 300

• nums[i] is 0, 1, or 2.

Input: nums = [2,0,2,1,1,0]
Output: [0,0,1,1,2,2]

Input: nums = [2,0,1]
Output: [0,1,2]

Input: nums = [0]
Output: [0]

Input: nums = [1]
Output: [1]

💡 Ideas

To achieve one-pass using only constant extra space, we use three-pointers instead of adopting count sort or recording any extra information.

• Pointer 0 for placing number 0

• Pointer 2 for placing number 2

• Pointer move for iterating the input array and placing the number 1

If you would like to do some optimization, we could find the consecutive 0 or 2 from the beginning and the end respectively.

In this way, we could check whether the input array is completely filled with 0 or 2 and shorten the runtime for the core algorithm.

🤖 Python3

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        # (base case)
        if len(nums) == 1: return
        
        # ==================================================
        #  Array + Three Pointers                          =
        # ==================================================
        # time  : O(n)
        # space : O(1)
        
        p0, p2 = 0, len(nums) - 1
        
        # (optimization)
        while p0 < len(nums) and nums[p0] == 0: p0 += 1
        if p0 == len(nums): return 
        
        while p2 >= 0 and nums[p2] == 2: p2 -= 1
        if p2 == 0: return
        
        moveP = p0
        
        while moveP <= p2:
            if nums[moveP] == 0:
                nums[p0], nums[moveP] = nums[moveP], nums[p0]
                p0 += 1
                moveP += 1
                
            elif nums[moveP] == 1:
                moveP += 1
                
            elif nums[moveP] == 2:
                nums[p2], nums[moveP] = nums[moveP], nums[p2]
                p2 -= 1

🤖 Java

class Solution {
    /**
     * @time  : O(n)
     * @space : O(1)
     */

    public void sortColors(int[] nums) {
        /* base case */
        if(nums.length == 1) return;
        
        int p0 = 0, p2 = nums.length - 1;
        
        while(p0 < p2 && nums[p0] == 0) p0++;
        if(p0 == p2) return;
            
        while(p2 <= 0 && nums[p2] == 2) p2--;
        if(p2 == 0) return;
            
        int moveP = p0;
        
        while(moveP <= p2) {
            if(nums[moveP] == 0) {
                int tmp = nums[p0];
                nums[p0] = nums[moveP];
                nums[moveP] = tmp;
                
                moveP++;
                p0++;
            } else if(nums[moveP] == 2) {
                int tmp = nums[p2];
                nums[p2] = nums[moveP];
                nums[moveP] = tmp;
                
                p2--;
            } else {
                moveP++;
            }
        }
    }
}