import random

def partition(arr, low, high, pivot_index):
    """Lomuto partition scheme moving elements smaller than pivot to the left."""
    pivot_value = arr[pivot_index]
    # Move pivot to the end
    arr[pivot_index], arr[high] = arr[high], arr[pivot_index]
    store_index = low

    for i in range(low, high):
        if arr[i] < pivot_value:
            arr[i], arr[store_index] = arr[store_index], arr[i]
            store_index += 1

    # Move pivot to its final sorted position
    arr[store_index], arr[high] = arr[high], arr[store_index]
    return store_index


def quickselect(arr, k_smallest, low=0, high=None):
    """Finds the k-th smallest element (1-indexed) using Randomized QuickSelect."""
    arr = arr.copy()
    if high is None:
        high = len(arr) - 1

    target_idx = k_smallest - 1  # Convert to 0-indexed position

    while low <= high:
        if low == high:
            return arr[low]

        # Choose a random pivot to maintain average O(N) performance
        pivot_index = random.randint(low, high)
        pivot_index = partition(arr, low, high, pivot_index)

        if target_idx == pivot_index:
            return arr[pivot_index]
        elif target_idx < pivot_index:
            high = pivot_index - 1
        else:
            low = pivot_index + 1


def find_kth_element(arr, k, find_largest=False):
    """Wrapper function to handle k-th smallest or k-th largest requests."""
    n = len(arr)
    if k < 1 or k > n:
        raise ValueError(f"k must be between 1 and {n}")

    if find_largest:
        # The k-th largest element is equivalent to the (n - k + 1)-th smallest
        target_k = n - k + 1
    else:
        target_k = k

    return quickselect(arr, target_k)


def get_median_of_five(group):
    """Sorts small groups (<= 5 elements) and returns the median."""
    group.sort()
    return group[len(group) // 2]


def select_pivot(arr, low, high):
    """Finds the Median-of-Medians pivot deterministically."""
    num_elements = high - low + 1
    if num_elements <= 5:
        return arr[low:high + 1].index(get_median_of_five(arr[low:high + 1])) + low

    # Step 1 & 2: Divide array into groups of 5 and find their medians
    medians = []
    for i in range(low, high + 1, 5):
        group = arr[i:min(i + 5, high + 1)]
        medians.append(get_median_of_five(group))

    # Step 3: Recursively find the median of medians
    median_of_medians_val = median_of_medians_select(
        medians, 0, len(medians) - 1, (len(medians) // 2) + 1
    )

    # Return the index of the pivot value in the original array slice
    for i in range(low, high + 1):
        if arr[i] == median_of_medians_val:
            return i


def median_of_medians_select(arr, low, high, k_smallest):
    """Deterministically finds the k-th smallest element in O(N) worst-case time."""
    if low == high:
        return arr[low]

    # Step 4: Select optimal pivot
    pivot_index = select_pivot(arr, low, high)
    pivot_index = partition(arr, low, high, pivot_index)

    # Step 5: Recurse based on target position
    target_idx = low + k_smallest - 1

    if target_idx == pivot_index:
        return arr[pivot_index]
    elif target_idx < pivot_index:
        return median_of_medians_select(arr, low, pivot_index - 1, k_smallest)
    else:
        # Adjust k relative to the new sub-array
        new_k = target_idx - pivot_index
        return median_of_medians_select(arr, pivot_index + 1, high, new_k)

if __name__ == "__main__":
    sample_data = [12, 3, 5, 7, 4, 19, 26, 23, 15, 8, 10, 2]
    # Sorted reference: [2, 3, 4, 5, 7, 8, 10, 12, 15, 19, 23, 26]
    
    k = 4
    
    # 1. QuickSelect Test
    kth_smallest = find_kth_element(sample_data, k=k, find_largest=False)
    kth_largest = find_kth_element(sample_data, k=k, find_largest=True)
    
    # 2. Median-of-Medians Test
    data_copy = sample_data.copy()
    mom_kth_smallest = median_of_medians_select(data_copy, 0, len(data_copy) - 1, k)

    print(f"Original Array: {sample_data}")
    print(f"Sorted Array:   {sorted(sample_data)}\n")
    print(f"{k}-th Smallest (QuickSelect):        {kth_smallest}")
    print(f"{k}-th Largest  (QuickSelect):        {kth_largest}")
    print(f"{k}-th Smallest (Median-of-Medians): {mom_kth_smallest}")
    