Problem Statement in English

You’re given a list of integers nums. Each number in nums can either be 00, 11, or 22, where each unique number represents a unique color.

Your task is to sort these numbers in constant time.

Approach

We use a sorting technique called bucket sort here. Let’s have a quick look at it.

Bucket sort is a sorting algorithm that works by distributing the elements of an array into a number of buckets.

In our case, since we know that there are only three possible colors (0, 1, and 2), we create three buckets. We then populate each bucket with the total occurrences of each color in the given list of integers.

Next, we overwrite the original list by using the count in each bucket. We start from the first bucket and iterate through each bucket, overwriting the original list with the corresponding color value. This process is done in-place, as required by the problem statement.

It’s important to note that we only apply this algorithm when we know that there’s a small number of buckets, otherwise this turns into an n2n^2 operation.

Solution in Python



class Solution:
    def sortColors(self, nums: List[int]) -> None:
        # we know that there are only 3 possible colors
        # so we create 3 buckets
        buckets = [0, 0, 0]

        # populate each bucket with total occurences in `nums`
        for val in nums:
            buckets[val] += 1

        # the current index in `nums` that we're overwriting at
        # since the question wants this to be in-place
        currentIndex = 0

        # use the count in each bucket to overwrite `nums`
        for bucketIndex in range(len(buckets)):
            while buckets[bucketIndex] > 0:
                nums[currentIndex] = bucketIndex
                buckets[bucketIndex] -= 1
                currentIndex += 1

Complexity

  • Time: O(n)O(n)
    Since we process at most nn elements in nums.

  • Space: O(1)O(1)
    Since the buckets variable stores only 3 elements, it’s considered constant space.

And we are done.