矢量化分组求和

Question

我有一个简单的Python循环算法，我想优化 速度。

从指示bin索引的数组开始，我想要为这些bin添加值。

更详细：我从一个指向值向量的索引数组开始，应该总结相同索引的值。普通和慢速Python版本就像：

import numpy

ix     = numpy.array([0 , 1 , 1 , 4 ])
values = numpy.array([10, 20, 30, 40])

# this models bin assignment:
# 10 belongs to bin 0
# 20 and 30 belong  to bin 1
# 40 belongs to bin 4

summed = numpy.zeros_like(values)

for i in ix:
    summed[i] += values[ix[i]]

print summed

[10, 50, 0, 0, 40]

这很慢，我问是否有人可以给我一个提示如何 矢量化。

Answer 1

您可以使用numpy.bincount()：

>>> numpy.bincount(ix, values)
array([ 10.,  50.,   0.,   0.,  40.])

Answer 2

有点破解，但这会得到相同的结果

indicator = np.arange(5)[None, ...] == ix[..., None]
summed = np.sum(values[..., None] * indicator, axis=0)