将多个值添加到一个numpy数组索引

时间:2017-02-01 21:25:02

标签: python arrays numpy

简单版本: 如果我这样做:

import numpy as np
a = np.zeros(2)
a[[1, 1]] += np.array([1, 1])

我将[0, 1]作为输出。但我想[0, 2]。这可能是某种方式,使用隐式numpy循环而不是自己循环它?

我真正需要做什么版本:

我有一个包含索引,值和一些布尔值的结构化数组。我想基于布尔值在那些索引处对这些值求和。显然,这可以通过一个简单的循环完成,但似乎应该可以通过巧妙的numpy索引(如上所述)。

例如,我有一个包含5个元素的数组,我想从数组中填充值,索引和条件:

import numpy as np
size = 5
nvalues = 10
np.random.seed(1)
a = np.zeros(nvalues, dtype=[('val', float), ('ix', int), ('cond', bool)])
a = np.rec.array(a)
a.val = np.random.rand(nvalues)
a.cond = (np.random.rand(nvalues) > 0.3)
a.ix = np.random.randint(size, size=nvalues)

# obvious solution
obvssum = np.zeros(size)
for i in a:
    if i.cond:
        obvssum[i.ix] += i.val

# is something this possible?
doesntwork = np.zeros(size)
doesntwork[a[a.cond].ix] += a[a.cond].val

print(doesntwork)
print(obvssum)

输出:

[ 0.          0.          0.61927097  0.02592623  0.29965467]
[ 0.          0.          1.05459336  0.02592623  1.27063303]

我认为这里发生的事情是a[a.cond].ix如果保证是唯一的,我的方法会正常工作,如简单示例中所述。

1 个答案:

答案 0 :(得分:7)

这是NumPy ufuncs的at方法的用途:

output = numpy.zeros(size)
numpy.add.at(output, a[a.cond].ix, a[a.cond].val)