Посчитай скользящее среднее окном k

Скользящее среднее — задача с алгоритмической секции уровня middle, и её ценность в том, что наивное решение пишется мгновенно, но оно медленное. Собеседующий ждёт, что ты сначала напишешь очевидный вариант, а потом сам увидишь, где теряется время, и ускоришь до линейного.

Задача

Дан массив чисел и размер окна k. Нужно вернуть массив средних для каждого окна длины k, скользящего по исходному массиву слева направо.

Для arr = [1, 2, 3, 4, 5, 6, 7] и k = 4 окон четыре: [1,2,3,4], [2,3,4,5], [3,4,5,6], [4,5,6,7] — их средние и есть ответ.

Решение

def moving_avg_naive(arr, k):
    result = []
    for i in range(len(arr) - k + 1):
        window = arr[i:i + k]
        result.append(sum(window) / k)
    return result


def moving_avg_window(arr, k):
    current = sum(arr[:k])
    result = [current / k]
    for i in range(len(arr) - k):
        current += arr[i + k] - arr[i]
        result.append(current / k)
    return result


arr = [1, 2, 3, 4, 5, 6, 7]
k = 4
naive = moving_avg_naive(arr, k)
window = moving_avg_window(arr, k)
print(f"наивное O(n*k):  {naive}")
print(f"скользящее O(n): {window}")
print(f"совпадают: {naive == window}")

Разбор

Начнём с наивного решения, потому что именно его пишут первым. Для каждой позиции окна берём срез arr[i:i+k], суммируем и делим на k. Работает, ответ верный. Но заметь: на каждом окне мы заново складываем k чисел, а окон примерно n. Итого O(n·k) — на длинном ряду и большом окне это медленно, и собеседующий обязательно спросит «а быстрее?».

Ускорение — приём скользящего окна. Соседние окна отличаются всего на два элемента: одно число ушло слева, одно пришло справа, а середина общая. Значит, пересчитывать всю сумму заново не нужно — достаточно из прошлой суммы вычесть уходящий элемент arr[i] и прибавить входящий arr[i + k]. Одно вычитание и одно сложение на шаг вместо k сложений. Сумма первого окна считается один раз, дальше катим её по массиву — получается O(n), один проход.

Две ловушки на краях, обе про счёт окон. Первая — сколько всего окон. Их len(arr) - k + 1, а не len(arr) - k. Легко ошибиться на единицу и потерять последнее окно [4,5,6,7]: в наивном варианте это выглядит как range(len(arr) - k), и ответ молча оказывается короче на одну строку. В решении наивная и оконная версии специально сверяются через == — это дешёвая проверка, что обе дают одинаковый результат и ни одно окно не потеряно.

Вторая — начальная сумма. В оконной версии первое среднее считается ДО цикла (sum(arr[:k])), а цикл добавляет только последующие окна, сдвигая сумму. Если забыть про начальное окно, потеряешь первую точку.

Итог: наивно O(n·k), окном O(n) — и оба на нашем примере дают [2.5, 3.5, 4.5, 5.5]. Тот же навык «не пересчитывай то, что уже посчитал» лежит в основе задачи про два указателя, где ответ тоже находят за один проход. А привычку ломать «очевидное» решение контрпримером хорошо тренирует задача про валидные скобки.

Ожидаемый результат

наивное O(n*k):  [2.5, 3.5, 4.5, 5.5]
скользящее O(n): [2.5, 3.5, 4.5, 5.5]
совпадают: True