Векторизованный минимум по корзинам в NumPy с reduceat

Как вычислить минимумы по корзинам в одномерном массиве NumPy без циклов и заполнения: используем np.minimum.reduceat для сегментной редукции по срезам.

Минимумы по корзинам в NumPy без циклов и без заполнения

Работа с одномерным массивом, разбитым на m корзин неравной длины, — распространённый сценарий. Каждый элемент связан с индексом корзины, который возрастает от 0 до m−1, а задача — вычислить минимум в каждой корзине. Прямолинейные подходы обычно либо выполняют цикл в Python, либо создают промежуточные структуры с заполнением. Оба варианта добавляют накладные расходы, которые плохо масштабируются.

Постановка задачи и типичный подход

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

import numpy as np

# данные: одномерный вектор
data_vec = np.array([...])  # длина n

# group_count = m, а group_slices — список пар (lo, hi)
# так что данные каждой корзины — это data_vec[lo:hi]

group_count = m
group_slices = [(lo0, hi0), (lo1, hi1), ...]  # неодинаковые длины

min_per_group = np.zeros(group_count)
for pos, (lo, hi) in enumerate(group_slices):
    min_per_group[pos] = np.min(data_vec[lo:hi])

Это работает, но цикл выполняется в пространстве Python и платит цену за каждую корзину. Альтернатива — дополнить до прямоугольника, сделать reshape и сократить по оси, но для этого приходится придумывать фиктивные элементы и трогать память, которая вам не нужна.

import numpy as np

# data_vec как выше
# valid_mask помечает реальные элементы среди заполненных ячеек
# scratch_1d заранее заполнен большим сторожевым значением; общий размер совпадает с «прямоугольником» после дополнения

scratch_1d[valid_mask] = data_vec
scratch_2d = scratch_1d.reshape((group_count, slot_size))
min_per_group = np.min(scratch_2d, axis=1)

В обоих вариантах либо есть цикл, либо приходится тасовать искусственные значения — это лишние проходы по данным, не влияющие на результат.

Почему так происходит

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

Векторизованное решение с np.minimum.reduceat

NumPy позволяет выполнить сегментную редукцию напрямую через np.minimum.reduceat. Единственное требование — массив стартовых смещений для каждой корзины длины m, где первый старт равен 0. С этими смещениями редукция один раз проходит исходный массив и вычисляет минимумы по сегментам без циклов и без дополнения.

import numpy as np

# data_vec: исходный одномерный массив (длина n)
# start_idx: одномерный массив длины m с начальным индексом каждой корзины; start_idx[0] должен быть равен 0

min_per_group = np.minimum.reduceat(data_vec, start_idx)

Здесь используется то, что индексы корзин возрастают от 0 до m−1, поэтому элементы каждой корзины образуют непрерывные срезы. Редукция запускается в каждой позиции из start_idx и векторизованно вычисляет минимумы по сегментам.

Зачем это знать

Такой подход избавляет от циклов в Python и не требует выдумывать фиктивные элементы. Память трогается ровно столько, сколько нужно; код упрощается и остаётся в оптимизированном пути выполнения NumPy. Для больших n и множества неравных корзин эта разница часто определяет, будет агрегация быстрой или медленной.

Выводы

Если ваш одномерный массив разбит на непрерывные корзины с индексами от 0 до m−1, храните массив начал корзин и используйте np.minimum.reduceat для вычисления минимумов по корзинам. Это соответствует уже имеющемуся расположению данных, сохраняет векторизацию и избавляет от дополнения и лишних проходов.