An implementation of the 'QuickSelect' algorithm (a.k.a. Hoare's selection algorithm)

quickselect_algorithm.py

def partition(data, left, right, pivot_idx):
    pivot_value = data[pivot_idx]
    real_pivot_idx = left
    data[pivot_idx], data[right] = data[right], data[pivot_idx]
    for i in range(left, right + 1):
        if data[i] < pivot_value:
            data[i], data[real_pivot_idx] = data[real_pivot_idx], data[i]
            real_pivot_idx += 1
    data[right], data[real_pivot_idx] = data[real_pivot_idx], data[right]
    return real_pivot_idx


def quickselect(data, n):
    left = 0
    right = len(data) - 1
    while True:
        if left == right:
            return data[left]
        pivot_idx = (left + right) / 2
        pivot_idx = partition(data, left, right, pivot_idx)
        if n == pivot_idx:
            return data[n]
        elif n < pivot_idx:
            right = pivot_idx - 1
        else:
            left = pivot_idx + 1
            
def median(data):
    m = quickselect(data, len(data) / 2)
    if len(data) % 2 == 0:
        m2 = quickselect(data, len(data) / 2 - 1)
        return (m2 + m) / 2.0
    else:
        return m