medians.py

"""
предполагаем что в последовательности нечётное число элементов для простоты
предполагаем что число элементов делится на 5 для простоты
"""


def pick_pivot(l):
    """
    Выбираем хорошй pivot в списке чисел l
    Этот алгоритм выполняется за время O(n).
    """

    # Если элементов < 5, просто возвращаем медиану
    if len(l) < 5:
        return nlogn_median(l)

    chunks = chunked(l, 5)

    sorted_groups = [sorted(chunk) for chunk in chunks]

    # Медиана каждого фрагмента имеет индекс 2
    medians = [chunk[2] for chunk in sorted_groups]

    # Мы находим медиану списка длиной n/5, поэтому эта операция также O(n)
    # Мы передаём нашу текущую функцию pick_pivot в качестве создателя pivot алгоритму
    # quickselect. O(n)
    median_of_medians = quickselect_median(medians, pick_pivot)
    return median_of_medians


def chunked(l, chunk_size):
    """Разделяем список `l` на фрагменты размером `chunk_size`."""
    return [l[i:i + chunk_size] for i in range(0, len(l), chunk_size)]


def quickselect_median(l, pivot_fn):
    return quickselect(l, len(l) / 2, pivot_fn)


def nlogn_median(l):
    l = sorted(l)
    return l[len(l) / 2]


def quickselect(l, k, pivot_fn):
    """
    :param pivot_fn: функция выбора pivot
    :return: k-тый по порядку элемент l
    """
    if len(l) == 1:
        assert k == 0
        return l[0]

    pivot = pivot_fn(l)

    lows   = [el for el in l if el < pivot]
    pivots = [el for el in l if el == pivot]
    highs  = [el for el in l if el > pivot]

    if k < len(lows):
        return quickselect(lows, k, pivot_fn)
    elif k < len(lows) + len(pivots):
        return pivots[0]
    else:
        k -= len(lows) + len(pivots)
        return quickselect(highs, k, pivot_fn)

предпологаем