Add 'Sort Colors'

Author: nickolashkraus

README.md

@@ -90,6 +90,8 @@ A collection of LeetCode solutions
 
 [Top K Frequent Elements](./src/top_k_frequent_elements.py)
 
+[Sort Colors](./src/sort_colors.py)
+
 [Two Sum](./src/two_sum.py)
 
 [Valid Anagram](./src/valid_anagram.py)

src/sort_colors.py

@@ -0,0 +1,27 @@
+"""
+75. Sort Colors
+
+https://leetcode.com/problems/sort-colors
+
+NOTES
+  * Use three-way partitioning (Dutch national flag problem).
+"""
+
+
+class Solution:
+    def sortColors(self, nums: list[int]) -> None:
+        """
+        Do not return anything, modify nums in-place instead.
+        """
+        i, j, k = 0, 0, len(nums) - 1
+
+        while j <= k:
+            if nums[j] > 1:
+                nums[j], nums[k] = nums[k], nums[j]
+                k -= 1
+            elif nums[j] < 1:
+                nums[j], nums[i] = nums[i], nums[j]
+                i += 1
+                j += 1
+            else:
+                j += 1

tests/test_sort_colors.py

@@ -0,0 +1,23 @@
+"""
+75. Sort Colors
+
+https://leetcode.com/problems/sort-colors
+"""
+
+from unittest import TestCase
+
+from src.sort_colors import Solution
+
+
+class TestSolution(TestCase):
+    def test_1(self):
+        exp = [0, 0, 1, 1, 2, 2]
+        nums = [2, 0, 2, 1, 1, 0]
+        Solution().sortColors(nums)
+        assert nums == exp
+
+    def test_2(self):
+        exp = [0, 1, 2]
+        nums = [2, 0, 1]
+        Solution().sortColors(nums)
+        assert nums == exp