[Leetcode problems]75. Sort Colors

Given an array 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.
Note: You are not suppose to use the library's sort function for this problem.
Input: [2,0,2,1,1,0] Output: [0,0,1,1,2,2]

This is known as a Dutch national flag problem – there are three colours in the flag of Netherlands: red, white, and blue, and given balls in these colours, the task is to sort these balls in the order of red, white, and blue. (Wikipedia) One of the potential applications of this problem is to develop a variant of partition()function in the quicksort algorithm.

Flag of Netherlands

The brute force solution could using a two-pass algorithm: use the first scanning to count the number of each colour and, for the second iterations, overwrite the colours according to their counts.

However, we are able to solve it with one-pass. We can initiate three pointers with i indicating the beginning of white balls, cur indicating the current position in the iteration, and j referring to the end of white balls. While iterating, once cur comes across a red ball, we can swap it with i, and i+=1; or if cur comes across a blue ball, we can swap it with j, and j-=1; or if cur comes across a white ball, it would simply move forward.

class Solution:
    def sortColors(self, nums: List[int]) -> None:
        i, cur, j = 0,0,len(nums)-1
        while cur <= j:
            if nums[cur] == 0:
                nums[i],nums[cur] = nums[cur],nums[i]
                cur += 1
                i += 1
            elif nums[cur] == 2:
                nums[j],nums[cur] = nums[cur],nums[j]
                j -= 1
                cur += 1

To provide a better explanation, I will also attach its pseudocode from Wikipedia:

procedure three-way-partition(A : array of values, mid : value):
    i ← 0
    j ← 0
    k ← size of A

    while j < k:
        if A[j] < mid:
            swap A[i] and A[j]
            i ← i + 1
            j ← j + 1
        else if A[j] > mid:
            k ← k - 1
            swap A[j] and A[k]
            j ← j + 1



Notify of
Inline Feedbacks
View all comments