Sort Colors

The problem description can be found in LeetCode.

Approach

A naive approach would consist of counting the three numbers in a static array of length 3.

nums = [2, 0, 2, 1, 1]

Would be:

counter = [1, 2, 2]
           |  |  |
           0  1  2

Where counter maps the number of occurrences of the numbers 0, 1 and 2.

Then, we could rearrange the elements based on this counter array. This is like using Counting sort.

from typing import List


class Solution:
    def sortColors(self, nums: List[int]) -> None:
        counter = [0, 0, 0]
        for num in nums:
            counter[num] += 1

        for i in range(len(nums)):
            if counter[0] > 0:
                counter[0] -= 1
                nums[i] = 0
            elif counter[1] > 0:
                counter[1] -= 1
                nums[i] = 1
            elif counter[2] > 0:
                counter[2] -= 1
                nums[i] = 2

Alternatively, we could think this problem as the Dutch national flag one.

Solution

from typing import List


class Solution:
    def sortColors(self, nums: List[int]) -> None:
        i = j = 0
        k = len(nums) - 1
        mid = 1

        while j <= k:
            if nums[j] < mid:
                nums[i], nums[j] = nums[j], nums[i]
                i += 1
                j += 1
            elif nums[j] > mid:
                nums[j], nums[k] = nums[k], nums[j]
                k -= 1
            else:
                j += 1

Complexity

• Time: $O(n)$
• Space: $O(1)$

Where $n$ is the size of the nums list.