使索引数组的循环更快

时间:2015-03-20 14:24:27

标签: python arrays numpy

我有以下问题:我有重复索引的索引数组,并希望像这样向数组添加值:

grid_array[xidx[:],yidx[:],zidx[:]] += data[:]

但是,由于我重复索引,因此它不能正常工作,因为numpy将创建一个临时数组,导致重复索引的数据被分配多次而不是相互添加(参见{{3} })。

这样的for循环
for i in range(0,n):
    grid_array[xidx[i],yidx[i],zidx[i]] += data[i]

会慢慢来。有没有办法我仍然可以使用numpy的矢量化?或者还有另一种方法可以更快地完成这项任务吗?

感谢您的帮助

4 个答案:

答案 0 :(得分:3)

如何使用bincount?

import numpy as np

flat_index = np.ravel_multi_index([xidx, yidx, zidx], grid_array.shape)
datasum = np.bincount(flat_index, data, minlength=grid_array.size)
grid_array += datasum.reshape(grid_array.shape)

答案 1 :(得分:2)

这是一个缓冲问题。 .at提供无缓冲操作 http://docs.scipy.org/doc/numpy/reference/generated/numpy.ufunc.at.html#numpy.ufunc.at

np.add.at(grid_array, (xidx,yidx,zidx),data)

答案 2 :(得分:0)

要将数组添加到嵌套数组的元素,您只需执行grid_array[::]+=data

>>> grid_array=np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> data=np.array([3,3,3])
>>> grid_array[::]+=data
>>> grid_array
array([[ 4,  5,  6],
       [ 7,  8,  9],
       [10, 11, 12]])

答案 3 :(得分:-1)

我想我找到了一个可能的解决方案:

def assign(xidx,yidx,zidx,data):
    grid_array[xidx,yidx,zidx] += data
    return

map(assign,xidx,yidx,zidx,sn.part0['mass'])