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

💡 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