k_th_statistic.py

def nth_element(array: List[Real], k: Real) -> Real:
    """
    Time Complexity: O(n)
    Space Complexity: O(1)

    >>> nth_element([1, 2, 3], 3)
    3
    >>> nth_element([-1, -2, 3], 2)
    -1
    >>> nth_element([-1, -2, 3, 10], 1)
    -2
    >>> nth_element([-1, -2, 3, 10], 3)
    3
    >>> nth_element([-1, -2, 3, 10], 4)
    10
    """
    return quick_select(array, 0, len(array) - 1, len(array) - k + 1)


def quick_select(nums, start, n, k):
    position = partition(nums, start, n)
    if position == k - 1:
        return nums[position]
    elif position >= k:
        return quick_select(nums, start, position - 1, k)
    return quick_select(nums, position + 1, n, k)


def partition(nums, left, right):
    pivot = nums[right]
    i = left
    for j in range(left, right):
        if nums[j] > pivot:
            nums[j], nums[i] = nums[i], nums[j]
            i += 1
    nums[right], nums[i] = nums[i], nums[right]
    return i