Sort Colors

main.py

class Solution:
    def sortColors(self, arr: List[int]) -> None:
        """
        Do not return anything, modify nums in-place instead.
        """
        # Move zeros to the first non-zero ind and twos to the
        # last non-two ind.
        n = len(arr)
        start, end = 0, n - 1
        ind = 0
        while ind <= end:
            # We incremenent ind when we move a 0 to the front but not
            # when we move a two to the back because since we are moving zeros
            # to the part of array which we have already iterated upon -- it is
            # garunteed that the zero will be swapped with a 1 or 0. But when we
            # swap a two to the back, the value swapped could be 0, 1, or 2, and
            # if it's a 0, we need to swap it again so we don't just incr ind
            # then and there.
            if arr[ind] == 0:
                arr[ind], arr[start] = arr[start], arr[ind]
                start += 1
                ind += 1
            elif arr[ind] == 2:
                arr[ind], arr[end] = arr[end], arr[ind]
                end -= 1
            else:
                ind += 1