如何用纯NumPy实现对可变大小分箱的最小值计算？

假设有一个大小为n的一维数组，每个元素对应一个从0到m-1的递增索引（m < n），示例如下：

[a1, a2, a3, a4, ..., aN]  # 原数组
[0,  1,  1,  1,  ..., m-1] # 关联索引

可将其视为被划分为m个分箱，分箱大小不一定相同。需求是得到一个大小为m的数组，包含每个分箱的最小值。

目前想到的方法要么用到非NumPy循环/列表，要么会产生大量冗余数据：

• 遍历分箱：

  mins = np.zeros(m)
  for i, (begin, end) in enumerate(bins):
      mins[i] = np.min(myarray[begin:end])
  

• 添加虚拟数据使分箱大小一致后重塑维度：假设已准备好大小为(m * bin_size)的mask（非虚拟位置为True），以及初始化为大常量的fake_array：

  fake_array[mask] = myarray
  fake_array_2d = fake_array.reshape((m, bin_size))
  mins = np.min(fake_array_2d, axis=1)
  

• 在(m, n)矩阵中计算（同样初始化为大数值）：

  np.put_along_axis(fake_matrix, indices[None, :], myarray, axis=0) # 调试了很久才弄对
  mins = np.min(fake_matrix, axis=1)
  

请问是否有更高效的纯NumPy实现方案？

可以利用np.minimum.reduceat实现，这是NumPy专门针对按分组索引进行归约计算的优化函数，完全不需要循环或虚拟数据，效率极高。

实现代码 #

import numpy as np

# 假设原数组为myarray，关联索引为indices（需满足题目中递增的条件）
# 1. 定位每个分箱的起始位置
split_indices = np.concatenate([[0], np.where(np.diff(indices) != 0)[0] + 1])

# 2. 按分组计算最小值
mins = np.minimum.reduceat(myarray, split_indices)

原理说明 #

1. np.diff(indices)：计算索引数组的相邻差值，差值不为0的位置对应新分箱的起始点前一位。
2. np.where(...) +1：将上述位置加1，得到每个新分箱在原数组中的起始索引。
3. np.concatenate([[0], ...])：补上第一个分箱的起始位置0，形成完整的分组起始索引列表。
4. np.minimum.reduceat：根据分组起始索引对原数组分段执行取最小值的归约操作，直接输出每个分箱的最小值数组。

优势对比 #

• 无循环依赖：完全依托NumPy底层优化，比Python循环效率高数个数量级。
• 零冗余数据：无需构造虚拟数组或矩阵，内存占用极低。
• 代码简洁：仅两行核心代码即可完成需求。

内容的提问来源于stack exchange，提问作者emengd