python - 多次更新python numpy数组列

Question

假设我有一个 2x3 矩阵 A:

1 2 3   
4 5 6  

和长度为 4 的向量 y:

0 1 2 1

以及另一个 4x2 矩阵 B:

0 0  
1 1   
2 2   
3 3 

我想通过添加 B 的行来多次更新 A 的列。A 要更新的列索引由 y 给出。

使用for循环，可以这样做:

for i in np.arange(4):
  A[:,y[i]] += B[i,:]

我使用ufunc.at实现了这个:

np.add.at(A.T,y,B)  

但是，ufunc.at 的性能几乎与使用 for 循环一样糟糕。

如何获得不同的矢量化实现？

使用 A[:,y]+=B.T 更新似乎每列仅更新一次。

Answer 1

方法#1

这是一种使用 np.add.reduceat 间隔求和的方法-

def reduceat_app(A, y, B):
    idx = y.argsort()
    y0 = y[idx]
    sep_idx = np.concatenate(([0], np.flatnonzero(y0[1:] != y0[:-1])+1))
    A += np.add.reduceat(B[idx],sep_idx).T

方法#2

由于A中的行数相对较少，我们还可以使用 np.bincount以迭代方式对每一行执行基于 bin 的求和，如下所示 -

def bincount_loopy_app(A, y, B): 
    n = A.shape[1]
    for i,a in enumerate(A):
        a += np.bincount(y,B[:,i],minlength=n).astype(A.dtype)

方法#3

我们可以通过为所有元素创建一个由 y 索引/箱组成的 2D 网格来矢量化以前的方法，这样对于每一行，我们都会有偏移的箱。通过该偏移，bincount 可用于以矢量化方式对所有行执行基于 bin 的求和。

因此，实现将是 -

def bincount_vectorized_app(A, y, B): 
    m,n = A.shape
    IDs = y[:,None] + n*np.arange(B.shape[1])
    vals = np.bincount(IDs.ravel(), B.ravel(), minlength=m*n)
    A += vals.astype(A.dtype).reshape((m,n))