[LeetCode] 75. Sort Colors

Sort Colors

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.

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]

Constraints:

n == nums.length
1 <= n <= 300

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

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

Solution

Plain solution: in-place quick sort.

To implement in-place quick sort, we need to recursively arrange subarrays so that the left part of each subarray is less than the pivot, while the right part is greater than the pivot. The overall time complexity is $O(n\log n)$.

class Solution(object):
    def quickSort(self, data, le, ri):
        if le >= ri:
            return
        pivotIndex = self.partition(data, le, ri)
        self.quickSort(data, le, pivotIndex-1)
        self.quickSort(data, pivotIndex+1, ri)
    
    def partition(self, data, le, ri):
        i = le-1
        for j in range(le, ri):
            if data[j] < data[ri]:
                i += 1
                data[i], data[j] = data[j], data[i]
        data[i+1], data[ri] = data[ri], data[i+1]
        return i+1

    def sortColors(self, nums):
        self.quickSort(nums, 0, len(nums)-1)
        

Referring to the above process, there are only three kinds of value $0,1,2$. So, the pivot must be $1$ and we just need to put $0$ to the left of the array and put $2$ to right of the array.

class Solution(object):
    def sortColors(self, nums):
        low = -1
        for i in range(len(nums)):
            if nums[i] == 0:
                low += 1
                nums[low], nums[i] = nums[i], nums[low]
        high = len(nums)
        for i in range(len(nums)-1, -1, -1):
            if nums[i] == 2:
                high -= 1
                nums[high], nums[i] = nums[i], nums[high]
        

Additionally, the two loops can be merged into one loop. We will put all $0$ at the left of $low$ and put all $2$ at the right of $high$.

class Solution(object):
    def sortColors(self, nums):
        low = -1
        high = len(nums)
        for i in range(len(nums)):
            while i < high and nums[i] == 2:
                high -= 1
                nums[high], nums[i] = nums[i], nums[high]
            while low < i and nums[i] == 0:
                low += 1
                nums[low], nums[i] = nums[i], nums[low]