diff --git a/medium/sort_colors.py b/medium/sort_colors.py new file mode 100644 index 0000000..9161116 --- /dev/null +++ b/medium/sort_colors.py @@ -0,0 +1,22 @@ +# https://leetcode.com/problems/sort-colors + +from typing import List + +class Solution: + def sortColors(self, nums: List[int]) -> None: + def swap(i: int, j: int): + nums[i], nums[j] = nums[j], nums[i] + i = 0 + j = 0 + k = len(nums) - 1 + while j <= k: + if nums[j] < 1: + swap(i, j) + i += 1 + j += 1 + elif nums[j] > 1: + swap(j, k) + k -= 1 + else: + j += 1 +