如何实现二维直方图对应np.digitize功能获取各分桶原始索引

实现方案 #

你可以通过分别对x、y坐标做bin映射，再按bin分组索引的方式实现需求，以下是可直接衔接你现有代码的实现：

基础实现（易理解，适合中小数据量）

import numpy as np

# --------------- 你的原有代码 ---------------
# 假设coord是你的原始坐标数组，shape为 (N, 2)
H, xedges, yedges = np.histogram2d(coord[:, 0], coord[:, 1])
H = H.T
# -------------------------------------------

# 计算每个点对应的x方向bin编号（从0开始，和H的列对应）
x_bin = np.digitize(coord[:, 0], xedges) - 1
# 适配histogram2d最右bin包含右边界的规则
x_bin[x_bin == len(xedges) - 1] = len(xedges) - 2
# 过滤x方向超出直方图范围的点
valid_x = (x_bin >= 0) & (x_bin < len(xedges) - 1)

# 计算每个点对应的y方向bin编号（从0开始，和H的行对应）
y_bin = np.digitize(coord[:, 1], yedges) - 1
y_bin[y_bin == len(yedges) - 1] = len(yedges) - 2
# 过滤y方向超出直方图范围的点
valid_y = (y_bin >= 0) & (y_bin < len(yedges) - 1)

# 提取所有在直方图范围内的有效点、对应原始索引、bin编号
valid_mask = valid_x & valid_y
valid_original_indices = np.arange(len(coord))[valid_mask]
valid_x_bin = x_bin[valid_mask]
valid_y_bin = y_bin[valid_mask]

# 初始化和H形状一致的索引存储结构
bin_indices = [[[] for _ in range(H.shape[1])] for _ in range(H.shape[0])]
# 把每个有效点的原始索引归入对应bin
for idx, yb, xb in zip(valid_original_indices, valid_y_bin, valid_x_bin):
    bin_indices[yb][xb].append(idx)

结果验证

你可以通过以下代码验证结果和直方图计数完全匹配：

for i in range(H.shape[0]):
    for j in range(H.shape[1]):
        assert len(bin_indices[i][j]) == int(H[i][j])
print("所有bin计数匹配验证通过")

你要的第一行第一列的7个点的原始索引就是bin_indices[0][0]。 #

优化实现（适合百万级以上大数据量）

如果数据量很大，Python级别的循环会比较慢，可以用numpy向量化分组提升效率：

# 将二维bin编码为一维值
bin_1d = valid_y_bin * H.shape[1] + valid_x_bin
# 按一维bin编号排序
sorted_order = bin_1d.argsort()
sorted_bin_1d = bin_1d[sorted_order]
sorted_indices = valid_original_indices[sorted_order]
# 找到不同bin的分界位置
split_pos = np.where(np.diff(sorted_bin_1d))[0] + 1
# 拆分得到每个bin的索引组
bin_groups = np.split(sorted_indices, split_pos)

# 映射回二维bin结构
bin_indices = [[[] for _ in range(H.shape[1])] for _ in range(H.shape[0])]
for b1d, group in zip(sorted_bin_1d[np.r_[0, split_pos]], bin_groups):
    yb = b1d // H.shape[1]
    xb = b1d % H.shape[1]
    bin_indices[yb][xb] = group.tolist()

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