如何利用np.hist分箱索引在Python中为其他数组分箱？

需求背景 #

使用np.hist(array, bins=100)完成数据分箱后，需要获取每个分箱内数据点的索引，进而从array1、array2等其他数组中提取对应值，得到100组数据。现有循环遍历每个区间、用np.where筛选索引的方法效率较低，需要更优方案。

高效解决方案（纯Numpy实现） #

利用np.digitize完成向量化分箱，避免逐区间的条件判断，大幅提升效率：

1. 获取每个元素的分箱编号
   np.digitize可以一次性给出数组中每个元素所属的分箱索引，注意和np.hist的区间规则对齐（默认左闭右开，最后一个区间为闭区间）：

   import numpy as np
   
   # 假设已得到分箱结果
   histo, edges = np.hist(array, bins=100)
   
   # 获取每个元素对应的分箱编号（1到100）
   bin_indices = np.digitize(array, edges, right=False)
   # 修正最后一个区间的边界值：np.hist最后一个区间包含edges[-1]，而digitize默认不包含，因此手动调整
   bin_indices[array == edges[-1]] = len(edges) - 1
   

2. 批量收集各分箱的索引
   基于分箱编号，一次性筛选出每个分箱对应的元素索引：

   bin_index_list = []
   # 遍历1到100号分箱
   for bin_num in range(1, len(edges)):
       # 直接匹配分箱编号，获取对应索引
       indices = np.where(bin_indices == bin_num)[0]
       bin_index_list.append(indices)
   

3. 提取其他数组的对应子集
   通过列表推导式快速从array1、array2中提取各分箱的对应数据：

   # 得到array1的100组分箱数据
   array1_groups = [array1[indices] for indices in bin_index_list]
   # 得到array2的100组分箱数据
   array2_groups = [array2[indices] for indices in bin_index_list]
   

可选：Pandas简化实现 #

如果允许使用Pandas，代码会更简洁，且分组逻辑更直观：

import pandas as pd

# 将数据整理为DataFrame
df = pd.DataFrame({
    'original': array,
    'array1': array1,
    'array2': array2
})

# 用pd.cut完成分箱，与np.hist的区间完全对齐
df['bin'] = pd.cut(df['original'], bins=edges, include_lowest=True)

# 按分箱分组，提取对应列的数据
groups = df.groupby('bin')
array1_groups = [group['array1'].values for _, group in groups]
array2_groups = [group['array2'].values for _, group in groups]

效率对比 #

原方案中每次循环都要执行(array>=edges[i])&(array<edges[i+1])的双条件判断，属于多次遍历数组；而np.digitize仅需一次向量化遍历即可完成所有元素的分箱标记，后续的索引筛选是基于整数匹配，速度提升显著，尤其当数组规模较大时优势更明显。

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