Sort Colors

Patrick Leaton
Problem Description

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.

Here, we will use the integers 01, and 2 to represent the color red, white, and blue respectively.

Follow up:

  • Could you solve this problem without using the library's sort function?
  • Could you come up with a one-pass algorithm using only O(1) constant space?

 

Example 1:

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

Example 2:

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

Example 3:

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

Example 4:

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

 

Constraints:

  • n == nums.length
  • 1 <= n <= 300
  • nums[i] is 01, or 2.

 

Description taken from https://leetcode.com/problems/sort-colors/.

Problem Solution

#O(N) Time, O(1) Space
class Solution:
    def sortColors(self, nums: List[int]) -> None:
        left, right = 0, len(nums)-1
        i, pivot01
        while i <= right:
            if nums[i] < pivot:
                nums[left], nums[i] = nums[i], nums[left]
                left += 1
                i += 
            elif nums[i] == pivot:
                i += 1
            else:
                nums[i], nums[right] = nums[right], nums[i]
                right -= 1

 

Problem Explanation


If we are given an array that we're required to sort in place, and we are given a predetermined middle value, our best bet is to solve this with a Quicksort-type algorithm.


Let's start by initializing a left pointer and an index pointer at the first index of the array, a right pointer at the last index, and a pivot value for what the middle values should be.  Here we want the pivot to be one so that we are swapping red and blue values around the white center so that we can sort our array into red, white, and blue.

        left, right = 0, len(nums)-1
        i, pivot01

 

Next, we will create our loop that will run while our index pointer hasn't crossed the right pointer.

We'll want to stop sorting once we have gone past the right pointer so that we don't begin swapping back values to their original place.

        while i <= right:

 

During each iteration, we will check for three cases.

The element in the current index could be less than the pivot, so it would be zero, red.

If this is the case, then we will swap the current value with the value at the left pointer.

            if nums[i] < pivot:
                nums[left], nums[i] = nums[i], nums[left]

 

After we have swapped with the left pointer, we will move our index pointer and our left pointer to the right.

Here we can imagine our left pointer walking up to the spot and staying there waiting for us to throw it a red piece of cloth in exchange for whatever is in the spot.   

                left += 1
                i += 1

 

The element in the current index could be equal to the pivot, so it would be one, white.

If this is the case, then we will keep moving forward.  The element should stay in the middle and we don't want to swap it left or right.

            elif nums[i] == pivot:
                i += 1

 

If the element in the current index is not less than or equal to the pivot, then it is greater so it would be two, blue.

If this is the case, we will swap our current element with the element at the right pointer.  There isn't a need to move the index pointer forward after this swap because we are already moving toward the right sorted partition since the right pointer is being decremented.

Like the left pointer, after our swap, we can imagine the right pointer walking down to a spot and stay there waiting for us to throw it a blue cloth in exchange for whatever is in the spot.  

            else:
                nums[i], nums[right] = nums[right], nums[i]
                right -= 1

 

We modified the array in place so we do not need to return anything.