Векторизованный подсчет длины возрастающих серий в NumPy

Узнайте, как в NumPy векторизованно считать длины возрастающих участков без циклов: маска роста, np.cumsum и maximum.accumulate, типы uint для ускорения.

Когда у вас есть одномерный массив NumPy, который в целом убывает, но местами подрастает, бывает полезно для каждого индекса получать длину текущего непрерывного участка роста. Трудность в том, что простой накопительный подсчет по булевой маске не сбрасывается, когда тренд снова идет вниз. Ниже приведен компактный, векторизованный способ вычислять длину каждого возрастающего участка с мгновенным сбросом счетчика, как только последовательность перестает расти. Результат сохраняет исходную форму массива, поэтому его легко визуализировать или обрабатывать дальше.

Пример, который проясняет задачу

Нам нужен массив той же длины, который для каждого индекса показывает, сколько подряд шагов вверх набрано в текущей серии; в остальных местах — нули.

import numpy as np
def rising_run_lengths(x: np.ndarray) -> np.ndarray:
    ...
arr = np.array([9, 8, 7, 9, 6, 5, 6, 7, 8, 4, 3, 1, 2, 3, 0])
res = rising_run_lengths(arr)
print(arr)
print(res)
# >>> [9 8 7 9 6 5 6 7 8 4 3 1 2 3 0]
# >>> [0 0 0 1 0 0 1 2 3 0 0 0 1 2 0]

Что происходит под капотом

Подумайте о булевой маске, которая помечает позиции, где значение массива больше предыдущего элемента. Если взять накопительную сумму этой маски, получится текущий счетчик увиденных повышений. Однако такой наивный счетчик никогда не сбрасывается — он растет на протяжении всего массива. Хитрость в том, чтобы для каждого возрастающего блока вычесть накопленное значение прямо перед началом этого блока. Чтобы сделать это векторизованно, берут накопительный максимум на позициях, где роста нет. Это дает последний «базовый уровень» для вычитания и тем самым точно сбрасывает накопление в момент, когда последовательность перестает расти.

Решение и рабочий код

Ниже реализовано вычисление булевой маски повышений, наивной накопительной суммы по этой маске, а затем удаление «залежавшегося» префикса через накопительный максимум, взятый по индексам без роста.

import numpy as np
def rising_run_lengths(x: np.ndarray) -> np.ndarray:
    up = np.diff(x, prepend=x[0]) > 0
    naive = np.cumsum(up)
    resets = np.maximum.accumulate(naive * ~up)
    return naive - resets
arr = np.array([9, 8, 7, 9, 6, 5, 6, 7, 8, 4, 3, 1, 2, 3, 0])
out = rising_run_lengths(arr)
print(arr)  # >>> [9 8 7 9 6 5 6 7 8 4 3 1 2 3 0]
print(out)  # >>> [0 0 0 1 0 0 1 2 3 0 0 0 1 2 0]

Идея лаконична и полностью векторизована. Сначала np.diff с параметром prepend строит маску «роста», выровненную с входом. Далее np.cumsum считает каждый True как один шаг вверх. Наконец, умножение текущего счетчика на отрицание маски сохраняет значения только в точках без роста, а np.maximum.accumulate протягивает последнее такое значение вперед; вычитание этого базового уровня обнуляет каждую новую серию роста.

Более читаемый вариант с np.where

Если хочется сделать логику сброса более явной, замените произведение на условный выбор. Функционально это тот же подход, просто выраженный по-другому.

import numpy as np
def rising_run_lengths_readable(x: np.ndarray) -> np.ndarray:
    up = np.diff(x, prepend=x[0]) > 0
    naive = np.cumsum(up)
    baseline = np.maximum.accumulate(np.where(up, 0, naive))
    return naive - baseline

Почему это полезно

Этот прием дает длину возрастающей серии в каждом индексе при сохранении формы входного массива — так проще строить графики и не приходится писать циклы. Он также демонстрирует переиспользуемую технику: построить накопительный счетчик и затем через накопительный максимум по замаскированному виду сбрасывать его ровно там, где нужно.

Можно явно указать целочисленные типы для шагов накопления. Сообщалось, что использование dtype=np.uint32 для накопительных операций на одной из машин ускорило вычисления примерно на 30%, и это безопасно для массивов меньше 4_000_000_000 элементов. Для небольших массивов, то есть меньше 65535 элементов, использование np.uint16 вместо этого, по сообщениям, давало ускорение порядка 60%.

import numpy as np
def rising_run_lengths_typed(x: np.ndarray) -> np.ndarray:
    up = np.diff(x, prepend=x[0]) > 0
    naive = np.cumsum(up, dtype=np.uint32)
    baseline = np.maximum.accumulate(np.where(up, 0, naive), dtype=np.uint32)
    return (naive - baseline).astype(np.uint32)

Вывод

Когда вам нужны длины непрерывных возрастающих подмассивов в NumPy, постройте булеву маску «роста», возьмите ее накопительную сумму и вычтите протянутый базовый уровень, полученный через накопительный максимум по позициям без роста. В итоге получается компактный, векторизованный результат, выровненный с исходным массивом, удобный для визуализации и последующего анализа. Если важна производительность, рассмотрите задание беззнакового целочисленного dtype для накопительных шагов — для указанных размеров это может дать заметный прирост скорости.