"""
Quickselect
-----------

Quickselect (also known as Hoare's selection algorithm) is a selection
algorithm to find the kth smallest (or largest) element in an unordered list of
n elements.

Quickselect uses the same overall approach as quicksort, choosing one element
as a pivot and partitioning the data in two based on the pivot, accordingly as
less than or greater than the pivot. However, instead of recursing into both
sides, as in quicksort, quickselect only recurses into one side-the side with
the element it is searching for. This reduces the average complexity from
O(nlog n) to O(n), with a worst case of O(n^2).

Why is quickselect O(n)?
------------------------
Since quickselect only cares about the partition which contains the element we
are looking for, half the elements are eliminated in each recursive call. This
forms a geometric series which sums to 2*n (or just n):

  n + n/2 + n/4 + ... ≈ 2n = O(n)
"""

from src.shared.partition import partition


def quickselect(l: list[int], left: int, right: int, k: int) -> int:
    """
    Return the kth element (0-based) in the given list.
    """
    if left == right:
        return l[left]

    # Retrieve the index of the pivot by partitioning the list into elements
    # less than and greater than or equal to the pivot.
    pivot = partition(l, left, right)

    # If k is equal to 'pivot', then l[pivot] is the kth element in the list.
    # Otherwise, execute quickselect on the partition comprising elements less
    # than or greater than or equal to the pivot. This partition is guaranteed
    # to contain the kth element.
    if k == pivot:
        return l[k]
    elif k < pivot:
        return quickselect(l, left, pivot - 1, k)
    else:
        return quickselect(l, pivot + 1, right, k)

"""
The partition function (in linear time) groups a list (ranging from indices
left to right) into two parts: those less than a certain element, and those
equal to or greater than the element.

This implementation uses the Lomuto partition scheme, which chooses as the
pivot the last element in the array.

  2 1 4 6 3 5
            ^
The algorithm maintains index i as it scans the array using another index j
such that the elements at lo through i-1 (inclusive) are less than the pivot,
and the elements at i through j (inclusive) are equal to or greater than the
pivot.

It is used in both the quickselect and quicksort algorithms.
"""


def partition(l: list[int], left: int, right: int) -> int:
    """
    Reorder the list such that elements less than the pivot are before elements
    equal to or greater than the pivot. When complete, the pivot is in its
    final sorted position. The pivot is chosen as the last element in the
    parition (Lomuto partition scheme).

    NOTE: We cannot simply use a slice of a list, since rearranging the
    elements in the slice will not reorder the elements in the original list.
    """

    def swap(l: list[int], i: int, j: int) -> None:
        """
        Swaps the element at index i with the element at index j.
        """
        temp = l[j]
        l[j] = l[i]
        l[i] = temp

    # Choose the last element (right) as the pivot.
    pivot = l[right]

    # i (commonly referred to as the "store index") is used to denote the index
    # of the pivot. j is used for scanning the list from left to right-1.
    i, j = left, left

    # The loop maintains the following invariant:
    #
    #   Elements left through i-1 (inclusive) are less than pivot
    #   Elements i through j (inclusive) are equal to or greater than pivot
    while j < right:
        if l[j] < pivot:
            swap(l, i, j)
            i += 1
        j += 1

    # As a final step, move pivot to its final position. This will be its final
    # position in the sorted array.
    swap(l, i, right)

    # Return the index of the pivot. The pivot index can be used to determine
    # the new left and right arguments for quicksort or can be used to
    # determine the kth element in a list.
    return i