Problem:

Given a list of numbers with only 3 unique numbers (1, 2, 3), sort the list in O(n) time.

Example 1:

Input: [3, 3, 2, 1, 3, 2, 1]
Output: [1, 1, 2, 2, 3, 3, 3]

def sortNums(nums):
  # Fill this in.

print sortNums([3, 3, 2, 1, 3, 2, 1])
# [1, 1, 2, 2, 3, 3, 3]

Challenge: Try sorting the list using constant space.

Solution 1:

Here we count the number of 1, 2, and 3, in the list, and then create a new list with the same numbers of 1, 2, and 3.

def sortNums(nums):
  num_one = 0
  num_two = 0
  num_three = 0
  
  for n in nums:
      if n == 1:
          num_one += 1
      elif n == 2:
          num_two += 1
      elif n == 3:
          num_three += 1
  
  return [1] * num_one + [2] * num_two + [3] * num_three

print sortNums([3, 3, 2, 1, 3, 2, 1])
# [1, 1, 2, 2, 3, 3, 3]

Solution 2:

We can sort the list in place if there are only 3 unique numbers. The general idea is to swap the lowest number to the leftmost index that is not sorted yet, and swap the highest number to the rightmost index that is not sorted yet.

In this case, we have one_idx that represents the leftmost index that is not sorted (so everything left of one_idx is a 1), and three_idx that represents the rightmost index that is not sorted( so everything right of three_idx is a 3).

We go through each element in the list and swap with one_idx if we see a 1, or swap with the three_idx if we see a 3. If the element is a 2 it is in the middle of the list already and no action needs to be done.

This code has a time complexity of O(n), as we only iterate through each element in the list once (once elements are swapped away we will never ‘see’ them again). The space complexity is O(1) as we are just swapping elements in the list.

def sortNums(nums):
  one_idx = 0 # Represents the leftmost index that is not 0
  three_idx = len(nums) - 1 # Represents the rightmost index that is not 2
  idx = 0 # Represents the index that we are checking
  while idx <= three_idx:
    if nums[idx] == 1:
      # Swap numbers with leftmost index
      # so all numbers up to one_idx are all 1
      nums[idx], nums[one_idx] = nums[one_idx], nums[idx]
      idx += 1
      one_idx += 1
    elif nums[idx] == 2:
      # If index is 2, it means list is in order, check next number
      idx += 1
    elif nums[idx] == 3:
      # Swap numbers with rightmost index
      # so all numbers after three_idx are all 3
      nums[idx], nums[three_idx] = nums[three_idx], nums[idx]
      three_idx -= 1
  return nums

print sortNums([3, 3, 2, 1, 3, 2, 1])
# [1, 1, 2, 2, 3, 3, 3]