使用两个一维数组高效地索引二维numpy数组

Question

我有一个大的2d numpy数组和两个表示2d数组中x/y索引的1d数组。我想使用这些1d数组对2d数组执行操作。使用for循环可以实现，但在处理大型数组时速度非常慢。是否有更快的方法？我尝试将1d数组简单地用作索引，但没有成功。请参考以下示例：

import numpy as np

# Two example 2d arrays
cnt_a   =   np.zeros((4,4))
cnt_b   =   np.zeros((4,4))

# 1d arrays holding x and y indices
xpos    =   [0,0,1,2,1,2,1,0,0,0,0,1,1,1,2,2,3]
ypos    =   [3,2,1,1,3,0,1,0,0,1,2,1,2,3,3,2,0]

# This method works, but is very slow for a large array
for i in range(0,len(xpos)):
    cnt_a[xpos[i],ypos[i]] = cnt_a[xpos[i],ypos[i]] + 1

# This method is fast, but gives incorrect answer
cnt_b[xpos,ypos] = cnt_b[xpos,ypos]+1


# Print the results
print 'Good:'
print cnt_a
print ''
print 'Bad:'
print cnt_b

这将输出以下内容：

Good:
[[ 2.  1.  2.  1.]
 [ 0.  3.  1.  2.]
 [ 1.  1.  1.  1.]
 [ 1.  0.  0.  0.]]

Bad:
[[ 1.  1.  1.  1.]
 [ 0.  1.  1.  1.]
 [ 1.  1.  1.  1.]
 [ 1.  0.  0.  0.]]

对于cnt_b数组，Numpy明显没有正确求和，但我不确定如何修复它，而不使用计算cnt_a时使用的（非常低效的）for循环。

- os1

---

你可以通过将循环中的行更改为 cnt_a[xpos[i],ypos[i]] += 1 来将第一个for循环的速度提高约一倍。 - Zinki

使用包含x和y的2列数组，这里有一个相关的Q&A。 - Divakar

Answer 1

另一种方法是使用1D索引（由@Shai建议），扩展以回答实际问题：

>>> out = np.zeros((4, 4))
>>> idx = np.ravel_multi_index((xpos, ypos), out.shape) # extract 1D indexes
>>> x = np.bincount(idx, minlength=out.size)
>>> out.flat += x

np.bincount 函数计算 xpos, ypos 中每个索引出现的次数，并将它们存储在 x 中。

或者，如 @Divakar 所建议的：

>>> out.flat += np.bincount(idx, minlength=out.size)

- Imanol Luengo

---

应该比 np.add.at 更快！ - Divakar

@Divakar 我不确定..我猜这取决于有多少索引和输出数组的大小。如果索引数量足够大，那么它可能会更快，但是如果目标数组很大，只有少数索引将被修改，那么np.add.at应该更快。 - Imanol Luengo

或者尝试这样做：out.ravel() = np.bincount(...)，省去一个步骤？需要帮助的话可以使用扁平化视图。 - Divakar

是的，由于某种奇怪的原因，“out.ravel() = ...”在Python3中无法工作（会出现“无法分配函数调用”的错误），但“out.flat = ...”应该可以正常工作。 - Imanol Luengo

Answer 2

我们可以计算线性索引，然后使用np.add.at将其累加到初始化为零的输出数组中。因此，假设有一个数组xpos和ypos，下面是一种实现方法 -

m,n = xpos.max()+1, ypos.max()+1
out = np.zeros((m,n),dtype=int)
np.add.at(out.ravel(), xpos*n+ypos, 1)

示例运行 -

In [95]: # 1d arrays holding x and y indices
    ...: xpos    =   np.array([0,0,1,2,1,2,1,0,0,0,0,1,1,1,2,2,3])
    ...: ypos    =   np.array([3,2,1,1,3,0,1,0,0,1,2,1,2,3,3,2,0])
    ...: 

In [96]: cnt_a   =   np.zeros((4,4))

In [97]: # This method works, but is very slow for a large array
    ...: for i in range(0,len(xpos)):
    ...:     cnt_a[xpos[i],ypos[i]] = cnt_a[xpos[i],ypos[i]] + 1
    ...:     

In [98]: m,n = xpos.max()+1, ypos.max()+1
    ...: out = np.zeros((m,n),dtype=int)
    ...: np.add.at(out.ravel(), xpos*n+ypos, 1)
    ...: 

In [99]: cnt_a
Out[99]: 
array([[ 2.,  1.,  2.,  1.],
       [ 0.,  3.,  1.,  2.],
       [ 1.,  1.,  1.,  1.],
       [ 1.,  0.,  0.,  0.]])

In [100]: out
Out[100]: 
array([[2, 1, 2, 1],
       [0, 3, 1, 2],
       [1, 1, 1, 1],
       [1, 0, 0, 0]])

- Divakar

---

谢谢，这比我的原始解决方案快一点。它还有一个优点，就是可以与原始数组的任何“添加”一起使用，而不仅仅是通过传递给np.add.at的最后一个变量添加1。 - os1

Answer 3

你可以同时迭代两个列表，并为每对元素递增（如果你不习惯，zip 可以将列表组合在一起）

for x, y in zip(xpos, ypos):
    cnt_b[x][y] += 1

但这将与您的解决方案A的速度相同。

如果您的xpos / ypos列表长度为n，我不认为您可以在少于 o（n）的时间内更新矩阵，因为您必须单向或双向检查每个对。

其他解决方案：您可以使用collections.Counter计算相似的索引对（例如：（0,3）等），并使用计数值更新矩阵。但我怀疑它不会更快，因为您在更新矩阵上节省的时间将在计算多个出现次数时丢失。

也许我完全错了，如果是这样，我也很想看到一个非 o（n）的答案

- smagnan

Answer 4

我认为你正在寻找ravel_multi_index函数

lidx = np.ravel_multi_index((xpos, ypos), cnt_a.shape)

将一维索引转换为cnt_a和cnt_b中的“平坦”索引：

np.add.at( cnt_b, lidx, 1 )

- Shai

---

我认为问题出在两次对相同坐标进行索引，而不是索引本身（例如，列表中的索引[0,0]出现了两次）。 - Imanol Luengo