python - 向量化分组求和

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)