如何高效提取Numpy二维直方图各分箱内的元素列表？

核心思路 #

利用分箱索引生成唯一分组键，结合numpy的向量化操作实现高效分组，避免Python层面的双重循环，大幅提升大数据集下的处理效率。

具体实现步骤 #

基于你已有的分箱索引代码，添加以下操作：

1. 预处理分箱索引

先将分箱索引转换为每个样本对应的(x分箱, y分箱)对：

import numpy as np

# 你的原有代码
bins = [[0, 0.3, 0.5, 0.7, 1.1], [0, 0.3, 0.7, 1.1]]
values = np.random.random((10, 2))
digitised = []
for i in range(len(bins)):
    digitised.append(np.digitize(values[:, i], bins[i], right=True))
digitised = np.concatenate(digitised).reshape(2, 10)

# 转换为每个样本的(x分箱, y分箱)数组，形状(样本数, 2)
sample_bins = digitised.T

2. 生成唯一分组键

将二维分箱索引转换为一维唯一整数键，确保每个(x,y)分箱对应唯一标识：

num_x_bins = len(bins[0]) - 1  # x方向分箱总数
num_y_bins = len(bins[1]) - 1  # y方向分箱总数

# 计算唯一分组键：x分箱编号 * y分箱总数 + y分箱编号
group_keys = sample_bins[:, 0] * num_y_bins + sample_bins[:, 1]

3. 高效分组元素

提供两种高效分组方式，按需选择：

方式一：排序拆分法（适合按分箱顺序获取元素）

# 按分组键排序样本
sorted_indices = np.argsort(group_keys)
sorted_values = values[sorted_indices]
sorted_keys = group_keys[sorted_indices]

# 找到键值变化的位置，拆分得到各分箱元素
split_points = np.where(np.diff(sorted_keys))[0] + 1
bin_elements = np.split(sorted_values, split_points)

# 映射到(x,y)分箱的字典
unique_keys = np.unique(group_keys)
bin_map = {}
for key, elem in zip(unique_keys, bin_elements):
    x_bin = key // num_y_bins
    y_bin = key % num_y_bins
    bin_map[(x_bin, y_bin)] = elem

方式二：唯一键映射法（直观对应分箱）

# 获取唯一分组键、逆索引
unique_keys, inverse_indices = np.unique(group_keys, return_inverse=True)

# 遍历生成(x,y)分箱到元素的映射
bin_map = {}
for idx, key in enumerate(unique_keys):
    # 筛选属于当前分箱的样本
    mask = inverse_indices == idx
    x_bin = key // num_y_bins
    y_bin = key % num_y_bins
    bin_map[(x_bin, y_bin)] = values[mask]

使用说明 #

• 最终bin_map是一个字典，键为(x分箱编号, y分箱编号)，值为对应分箱内的元素数组
• 对应你示例中的H[0][1]（y=0, x=1的分箱），可通过bin_map[(1, 0)]获取该分箱的所有元素
• 两种方法均基于numpy向量化操作，在8×7这类多分箱场景下，效率远高于双重循环判断

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