用 numpy 加速花式索引

Question

我有两个 numpy 数组，每个数组的形状为 (10000,10000)。 一个是值数组，一个是索引数组。

Value=np.random.rand(10000,10000)
Index=np.random.randint(0,1000,(10000,10000))

我想通过对引用“索引数组”的所有“值数组”求和来创建一个列表（或 1D numpy 数组）。例如，对于每个索引 i，找到匹配的数组索引并将其作为参数传递给值数组

for i in range(1000):
    NewArray[i] = np.sum(Value[np.where(Index==i)])

但是，这太慢了，因为我必须遍历 300,000 个数组。 我试图想出一些逻辑索引方法，例如

NewArray[Index] += Value[Index]

但它没有用。 我尝试的下一件事是使用字典

for k, v in list(zip(Index.flatten(),Value.flatten())):
    NewDict[k].append(v)

和

for i in NewDict:
    NewDict[i] = np.sum(NewDict[i])

但也很慢

有什么聪明的方法可以加快速度吗？

Answer 1

我有两个想法。首先，尝试使用遮罩，它的速度提高了大约 4 倍：

for i in range(1000):
    NewArray[i] = np.sum(Value[Index==i])

或者，您可以对数组进行排序，以将要相加的值放在连续的内存空间中。每次在切片上调用 sum 时，屏蔽或使用 where() 必须将所有值收集在一起。通过预先加载此聚会，您可能可以大大加快速度：

# flatten your arrays
vals = Value.ravel()
inds = Index.ravel()
s = np.argsort(inds)  # these are the indices that will sort your Index array

v_sorted = vals[s].copy()  # the copy here orders the values in memory instead of just providing a view
i_sorted = inds[s].copy()
searches = np.searchsorted(i_sorted, np.arange(0, i_sorted[-1] + 2)) # 1 greater than your max, this gives you your array end...
for i in range(len(searches) -1):
    st = searches[i]
    nd = searches[i+1]
    NewArray[i] = v_sorted[st:nd].sum()

此方法在我的计算机上需要 26 秒，而使用旧方法需要 400 秒。祝你好运。如果您想了解更多关于连续内存和性能的信息check this discussion out.