numpy中的矢量化聚集操作

时间:2016-09-18 11:11:18

标签: numpy vectorization

鉴于此(样本)数据

target_slots = np.array([1, 3, 1, 0, 8, 5, 8, 1, 1, 2])
dummy_elements = np.arange(10*D).reshape(10, D)

有没有办法在向量化的numpy表达式中表达操作

gathered_results = np.zeros((num_slots, D))
for i, target in enumerate(target_slots):
    gathered_results[target] += dummy_elements[i]

此操作看起来像bincount但不是计算我们对另一个数组的元素求和。

(暗示np.max(target_slots)<num_slotsnp.min(target_slots)>=0以及target_slots.shape[0] == D

1 个答案:

答案 0 :(得分:1)

方法#1

您正在执行间隔编辑求和选择dummy_elements之外的行并将特定行添加到输出数组中。因此,矢量化解决方案的一个明显选择是使用np.add.reduceat,就像这样 -

sidx = target_slots.argsort()
out = np.zeros((num_slots, D))
unq, shift_idx = np.unique(target_slots[sidx],return_index=True)
out[unq] = np.add.reduceat(dummy_elements[sidx],shift_idx)

方法#2

或者,我们也可以使用np.bincount来执行这些基于ID的求和操作。一种方法是使用循环沿着dummy_elements的列进行迭代,我认为当没有时,这将是有益的。这些柱的比例相对较小。实现看起来像这样 -

out = np.zeros((num_slots, D))
L = target_slots.size
for i in range(D):
    out[:,i] = np.bincount(target_slots,dummy_elements[:,i],minlength=L)

方法#3

同样的矢量化版本就像这样 -

L = target_slots.size
ids = (target_slots[:,None] + np.arange(D)*L).ravel('F') 
out = np.bincount(ids,dummy_elements.ravel('F'),minlength=L*D).reshape(D,-1).T