如何在 Python 中向量化增量

Question

我有一个二维数组，我有一些数字要添加到一些单元格中。我想对操作进行矢量化以节省时间。问题是当我需要向同一个单元格添加多个数字时。在这种情况下，矢量化代码仅添加最后一个。 'a' 是我的数组，'x' 和 'y' 是我要递增的单元格的坐标，而 'z' 包含我要添加的数字。

import numpy as np

a=np.zeros((4,4))
x=[1,2,1]
y=[0,1,0]
z=[2,3,1]
a[x,y]+=z
print(a)

如您所见，a[1,0] 应该递增两次：1 加 2，1 加 1。所以预期的数组应该是：

[[0. 0. 0. 0.]
 [3. 0. 0. 0.]
 [0. 3. 0. 0.]
 [0. 0. 0. 0.]]

但我得到了：

[[0. 0. 0. 0.]
 [1. 0. 0. 0.]
 [0. 3. 0. 0.]
 [0. 0. 0. 0.]]

用for循环很容易解决这个问题，但我想知道我是否可以正确地向量化这个操作。

Answer 1

为此使用np.add.at：

import numpy as np

a = np.zeros((4,4))
x = [1, 2, 1]
y = [0, 1, 0]
z = [2, 3, 1]
np.add.at(a, (x, y), z)
print(a)
# [[0. 0. 0. 0.]
#  [3. 0. 0. 0.]
#  [0. 3. 0. 0.]
#  [0. 0. 0. 0.]]

Answer 2

当你在做a[x,y]+=z时，我们可以将操作分解为：

a[1, 0], a[2, 1], a[1, 0] = [a[1, 0] + 2, a[2, 1] + 3, a[1, 0] + 1]
# Equivalent to :
a[1, 0] = 2
a[2, 1] = 3
a[1, 0] = 1

这就是它不起作用的原因。 但是，如果您使用每个维度的循环来增加数组，它应该可以工作

Answer 3

您可以创建一个大小为 3x4x4 的多维数组，然后将 z 添加到所有 3 个不同的维度，然后将它们全部相加

import numpy as np
x = [1,2,1]
y = [0,1,0]
z = [2,3,1]
a = np.zeros((3,4,4))
n = range(a.shape[0])
a[n,x,y] += z
print(sum(a))

这将导致

[[0. 0. 0. 0.]
 [3. 0. 0. 0.]
 [0. 3. 0. 0.]
 [0. 0. 0. 0.]]

Answer 4

方法 #1：基于 Bincount 的性能方法

我们可以使用np.bincount 进行高效的基于bin 的求和，并且基本上受到this post 的启发-

def accumulate_arr(x, y, z, out):
    # Get output array shape
    shp = out.shape

    # Get linear indices to be used as IDs with bincount
    lidx = np.ravel_multi_index((x,y),shp)
    # Or lidx = coords[0]*(coords[1].max()+1) + coords[1]

    # Accumulate arr with IDs from lidx
    out += np.bincount(lidx,z,minlength=out.size).reshape(out.shape)
    return out

如果您使用的是零初始化的输出数组，请将输出形状直接输入到函数中，并将 bincount 输出作为最终输出。

给定样本的输出 -

In [48]: accumulate_arr(x,y,z,a)
Out[48]: 
array([[0., 0., 0., 0.],
       [3., 0., 0., 0.],
       [0., 3., 0., 0.],
       [0., 0., 0., 0.]])

方法 #2：使用稀疏矩阵提高内存效率

In [54]: from scipy.sparse import coo_matrix

In [56]: coo_matrix((z,(x,y)), shape=(4,4)).toarray()
Out[56]: 
array([[0, 0, 0, 0],
       [3, 0, 0, 0],
       [0, 3, 0, 0],
       [0, 0, 0, 0]])

如果您对稀疏矩阵没问题，请跳过.toarray() 部分以获得内存高效的解决方案。