Dutch national flag problem: four colors - in Python

dutch_flag_four_colors.py

# (Variant for exercise 5.1 on EPI (Elements of Programming Interviews))
# The rationale behind it is to squeeze the forth color (middle right color) in between the middle left and right sub-arrays. It defines the colors as the algorithm progresses.
# It has a O(n) time complexity and O(1) space complexity

from typing import List

def dutch_variant_four_colors(array: List[int]) -> List[int]:
    left = array[0]
    mid_left = None
    right = None

    left_i = 0
    mid_left_i = 0
    mid_right_i = len(array)
    right_i = len(array)

    while mid_left_i < right_i and mid_left_i < mid_right_i:
        if (array[mid_left_i] == left):
            array[mid_left_i], array[left_i] = array[left_i], array[mid_left_i]
            mid_left_i += 1
            left_i += 1
        elif (right is None or array[mid_left_i] == right):
            right_i -= 1
            mid_right_i = right_i
            array[mid_left_i], array[right_i] = array[right_i], array[mid_left_i]
            right = array[right_i]
        else:  # it is a mid value
            if (mid_left is None):
                mid_left = array[mid_left_i]
            if (array[mid_left_i] == mid_left):
                mid_left_i += 1
            else:
                mid_right_i -= 1
                array[mid_left_i], array[mid_right_i] = array[mid_right_i], array[mid_left_i]

    return array


# Sample usages
print(dutch_variant_four_colors([1,1,3,3,4,3,2,2]))
print(dutch_variant_four_colors([1,2,3,4,2,3,1,3]))
print(dutch_variant_four_colors([1,2,3,4,4]))
print(dutch_variant_four_colors([0,1,2,5,5,2,2,0]))
print(dutch_variant_four_colors([1,0,3,0,5,5]))
print(dutch_variant_four_colors([1,2,3,2,5,1,1,3]))
print(dutch_variant_four_colors([5,1,2,3,2,1,1,3]))
print(dutch_variant_four_colors([3,2,5,1]))
print(dutch_variant_four_colors([3,2,1,5]))
print(dutch_variant_four_colors([3,2,1,3,5,2,2,1,2,3,5,3,2,1,3,5,3,3,2,2,2,5,5,5,3,3,2,5,3,1,2,3,2,1,3,2,1,1,2,3,2,3,2,1,2,3,2,1,2,3,2,2,2,2,2,3,3,3,1,2,2,1,1,2,3]))