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.])