Sort Colors

Sources: leetcode 75

Rating: medium

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 0, 1, and 2 to represent the color red, white, and blue respectively.

  • Solve without using a sort library.
  • Sort in single-pass algorithm using constant space.
def sortColors(nums: List[int]) -> None:
    # Do not return anything, modify nums in-place instead.

Solution

This is the so-called Dutch flag problem. The solution splits the array into three groupings, iterating through using three indices l, m, h, exchanging elements at the indices to ensure that each element ends up in the right groupings.

def sortColors(nums):
    l, m, h = 0, 0, len(nums) - 1
    while m <= h:
        # exchange elements

If the element at m is equal to the first value, then it must be exchanged for an element at l.

def sortColors(nums):
    l, m, h = 0, 0, len(nums) - 1
    while m <= h:
        if nums[m] == 0:
            nums[m], nums[l] = nums[l], nums[m]
            l += 1
            m += 1
        # exchange elements

If the element at m is the second value, then it is in the m grouping and the index advances.

def sortColors(nums):
    l, m, h = 0, 0, len(nums) - 1
    while m <= h:
        if nums[m] == 0:
            nums[m], nums[l] = nums[l], nums[m]
            l += 1
            m += 1
        elif nums[m] == 1:
            m += 1
        # exchange elements

Otherwise, m is the third value, and must be exchanged with an element in the h grouping, and h updates.

Final Code

def sortColors(nums):
    l, m, h = 0, 0, len(nums) - 1
    while m <= h:
        if nums[m] == 0:
            nums[m], nums[l] = nums[l], nums[m]
            l += 1
            m += 1
        elif nums[m] == 1:
            m += 1
        else:
            nums[m], nums[h] = nums[h], nums[m]
            h -= 1

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s