数组中的Numpy sum元素基于其值

时间:2015-11-20 02:41:25

标签: python arrays performance numpy

我有未排序的索引数组:

i = np.array([1,5,2,6,4,3,6,7,4,3,2])

我还有一组长度相同的值:

v = np.array([2,5,2,3,4,1,2,1,6,4,2])

我的数组中包含所需值的零:

d = np.zeros(10)

现在我想基于它在i中的索引添加v的d值中的元素。

如果我在普通的python中这样做,我会这样做:

for index,value in enumerate(v):
    idx = i[index]
    d[idx] += v[index]

这是丑陋而低效的。我怎样才能改变它?

2 个答案:

答案 0 :(得分:5)

np.add.at(d, i, v)

您认为d[i] += v会起作用,但如果您尝试以相同的方式对这些单元格进行多次添加,则其中一个会覆盖其他单元格。 ufunc.at方法避免了这些问题。

答案 1 :(得分:4)

我们可以使用np.bincount这对于这种累积加权计数非常有效,所以这里有一个 -

counts = np.bincount(i,v)
d[:counts.size] = counts

或者,使用minlength输入参数和d可以是任何数组并且我们想要添加到其中的一般情况 -

d += np.bincount(i,v,minlength=d.size).astype(d.dtype, copy=False)

运行时测试

本部分将other post中列出的基于np.add.at的方法与本文前面列出的基于np.bincount的方法进行了比较。

In [61]: def bincount_based(d,i,v):
    ...:     counts = np.bincount(i,v)
    ...:     d[:counts.size] = counts
    ...: 
    ...: def add_at_based(d,i,v):
    ...:     np.add.at(d, i, v)
    ...:     

In [62]: # Inputs (random numbers)
    ...: N = 10000
    ...: i = np.random.randint(0,1000,(N))
    ...: v = np.random.randint(0,1000,(N))
    ...: 
    ...: # Setup output arrays for two approaches
    ...: M = 12000
    ...: d1 = np.zeros(M)
    ...: d2 = np.zeros(M)
    ...: 

In [63]: bincount_based(d1,i,v) # Run approaches
    ...: add_at_based(d2,i,v)
    ...: 

In [64]: np.allclose(d1,d2)  # Verify outputs
Out[64]: True

In [67]: # Setup output arrays for two approaches again for timing
    ...: M = 12000
    ...: d1 = np.zeros(M)
    ...: d2 = np.zeros(M)
    ...: 

In [68]: %timeit add_at_based(d2,i,v)
1000 loops, best of 3: 1.83 ms per loop

In [69]: %timeit bincount_based(d1,i,v)
10000 loops, best of 3: 52.7 µs per loop