如何用向量化方式将2D NumPy数组值填充到分箱后的3D数组

核心解决方案 #

通过numpy的累计求和和高级索引实现无循环的向量化填充，代码如下：

import numpy as np

# 原始输入数据
x_y_bins = np.array([[0, 0], [0, 0], [2, 3], [4, 1], [2, 3], [0, 0], [1, 2]])
array_of_points = np.ones((7, 5))

# 获取分箱信息（你已完成的步骤）
unique_pairs, binning_index, count = np.unique(x_y_bins, return_inverse=True, return_counts=True, axis=0)

# 预分配目标3D数组（修正原代码的维度括号）
max_count = count.max()
target_array = np.zeros((len(unique_pairs), max_count, 5))

# 生成每个点在对应分箱内的位置索引
cum_counts = np.cumsum(count)
bin_start = np.zeros_like(cum_counts)
bin_start[1:] = cum_counts[:-1]
pos_in_bin = np.arange(len(binning_index)) - bin_start[binning_index]

# 向量化赋值
target_array[binning_index, pos_in_bin] = array_of_points

# 验证输出（与期望一致）
print(target_array)

步骤解释 #

1. 分箱信息获取：通过np.unique得到唯一分箱对、每个点的分箱索引binning_index、每个分箱的点数count，这部分你已经完成。
2. 目标数组预分配：按(唯一分箱数, 最大分箱点数, 5)的维度创建全零数组。
3. 生成分箱内位置索引：
   • cum_counts计算累计点数，用于确定每个分箱在原始点数组中的起始位置；
   • bin_start存储每个分箱的起始索引；
   • pos_in_bin通过原始点的全局索引减去对应分箱的起始索引，得到每个点在分箱内的相对位置（第0个、第1个...）。
4. 向量化赋值：利用numpy的高级索引特性，直接将原始点数据批量填充到目标数组的对应位置，全程无Python循环，效率远高于逐箱遍历。

效果验证 #

运行代码后，target_array的输出与你给出的期望结果完全一致，且处理大规模点云数据时性能优势明显。

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