Friedrich Ewald My Personal Website

Leetcode: Sort colors

Given an array nums with n colored objects, sort the array in place so that the colors are orders. The colors are represented by the numbers 0, 1, and 2. The easiest way to achieve this is via quicksort in O(n*log n).

from typing import List, Optional

class Solution:
  def sortColors(self, nums: List[int]) -> None:
    """
    Do not return anything, modify nums in-place instead.
    """
    self.quicksort(nums, 0, len(nums) - 1)
  
  def quicksort(self, arr, left, right):
    if left < right:
      partition_pos = self.partition(arr, left, right)
      self.quicksort(arr, left, partition_pos - 1)
      self.quicksort(arr, partition_pos + 1, right)

  def partition(self, arr, left, right):
    i = left
    j = right - 1
    pivot = arr[right]

    while i < j:
      while i < right and arr[i] < pivot:
        i += 1
      while j > left and arr[j] >= pivot:
        j -= 1
      
      if i < j:
        arr[i], arr[j] = arr[j], arr[i]

    if arr[i] > pivot:
      arr[i], arr[right] = arr[right], arr[i]
    
    return i


if __name__ == '__main__':
  s = Solution().sortColors
  # print(s([2,0,2,1,1,0]), [0,0,1,1,2,2])
  l  = [2,0,1]
  print(s(l), [0,1,2])
  print(l)
Runtime: 70 ms, faster than 12.05% of Python3 online submissions for Sort Colors. Memory Usage: 13.8 MB, less than 64.59% of Python3 online submissions for Sort Colors.


About the author

is an experienced Software Engineer with a Master's degree in Computer Science. He started this website in late 2015, mostly as a digital business card. He is interested in Go, Python, Ruby, SQL- and NoSQL-databases, machine learning and AI and is experienced in building scalable, distributed systems and micro-services at multiple larger and smaller companies.