"""
Problem:
    Write a program that takes an array A and an index i into A, and rearranges
    the elements such that all elements less than A[i] (the "pivot") appear
    first, followed by elements equal to the pivot, followed by elements
    greater than the pivot.

Hint:
    Think about the partition step in quicksort.

Constraints:
    xxx

Ideas:
    - using additional space O(n) where n is the length of A, we can have three
    lists namely elements less than the pivot, elements equal to the pivot and
    elements greater than the pivot. We then write these values into A.
    The time complexity is O(n).
    - Alternatively, avoid use of O(n) additional space at the cost of
    incurring more time by iterating through A and on each iteration, we find
    an element smaller than the pivot and move it into subarray of smaller
    elements via exchange. This moves all elements smaller than the pivot to
    the start of the array. In the second stage, we deal with the elements
    greater than the pivot in a similar fashion. These end up at the end of the
    array. See v1 solution.

Analysis:
    Space complexity is O(1) but time complexity is O(n^2).
"""


def dutch_flag(pivot_index: int, A: list):
    pivot = A[pivot_index]

    # group elements smaller than the pivot
    for i in range(len(A)):
        for j in range(i + 1, len(A)):
            if A[j] < pivot:
                A[i], A[j] = A[j], A[i]
                break

    # group elements larger than the pivot
    for i in reversed(range(len(A))):
        # if an element is less than pivot, stop since preceding
        # pass already moved them to start of A
        if A[i] < pivot:
            break
        for j in reversed(range(i)):
            if A[j] > pivot:
                A[i], A[j] = A[j], A[i]
                break

    return A


"""
Ideas:
    - Improving on v1 solution, we can do the operation in a single pass thus
    time complexity of O(n). In the first pass, we move all elements less than
    the pivot to the beginning and in the second pass, we move larger elements
    to the end. Each pass can be performed in a single iteration.

Analysis:
    Time complexity is O(2n) -> O(n) and space complexity is O(1).
"""


def dutch_flag_partition(pivot_index: int, A: list):
    pivot = A[pivot_index]

    # first pass
    smaller = 0
    for i in range(len(A)):
        if A[i] < pivot:
            A[i], A[smaller] = A[smaller], A[i]
            smaller += 1

    # second pass
    larger = len(A) - 1
    for i in reversed(range(len(A))):
        if A[i] < pivot:
            break
        elif A[i] > pivot:
            A[i], A[larger] = A[larger], A[i]
            larger -= 1

    return A


"""
Ideas:
    - Segregate elements into three buckets of less than, equal to and greater
    than pivot in a single pass.
    - the implemetation is tricker: maintain 4 subarrays -> bottom (less than
    pivot), middle (equal to pivot), top (greater than pivot) and unclassified.
    Move elements from unclassified to each of the 3 subarrrays.
    - Reduced runtime but trickier to implement. Each iteration reduces the
    size of unclassified by 1.

Analysis:
    - The time spent in each iteration is O(1) implying the time complexity is
    O(n). Space complexity is O(1).
"""


def dutch_flag_partition_2(pivot_index: int, A: list):
    pivot = A[pivot_index]
    smaller, equal, larger = 0, 0, len(A) - 1

    # keep the following invariants during partitioning:
    # bottom: A[:smaller]
    # middle: A[smaller:equal]
    # unclassified: A[equal:larger]
    # top: A[:larger]

    # keep iteration as long as there is an unclassified element
    while equal < larger:
        # A[equal] is the incoming unclassified element
        if A[equal] < pivot:
            A[smaller], A[equal] = A[equal], A[smaller]
            smaller, equal = smaller + 1, equal + 1
        elif A[equal] == pivot:
            equal += 1
        else:  # A[equal] > pivot
            A[equal], A[larger] = A[larger], A[equal]
            larger -= 1

    return A


"""
Variants of this problem:
    - Assuming keys take on of 3 values, reorder the array so that all elements
    with the same key appear together.
    - Given an array A of n objects with keys that take one of four values,
    reorder the array so that all objects that have the same key appear
    together. Use O(1) additional space and O(n) time.
    - Given an array A of n objects with boolean-valued keys, reorder the array
    so that objects that have key false appear first. The relative ordering
    of objects with key true should not change. Use additional space O(1) and
    time O(n).
"""
