python - 将给定点上的 3x3 数组求和到另一个保持边界的矩阵

Question

假设我有这个二维数组 A:

[[0,0,0,0],
 [0,0,0,0],
 [0,0,0,0],
 [0,0,0,4]]

我想对 B 求和:

[[1,2,3]
 [4,5,6]
 [7,8,9]]

以 A[0][0] 为中心，因此结果为:

array_sum(A,B,0,0) =
[[5,6,0,4],
 [8,9,0,0],
 [0,0,0,0],
 [2,0,0,5]]

我想我应该创建一个函数来比较它是否在边界上，然后调整索引:

def array_sum(A,B,i,f):
   ...
   if i == 0 and j == 0:
      A[-1][-1] = A[-1][-1]+B[0][0]
      ...

   else:
      A[i-1][j-1] = A[i][j]+B[0][0]
      A[i][j] = A[i][j]+B[1][1]
      A[i+1][j+1] = A[i][j]+B[2][2]
      ...

但我不知道是否有更好的方法来做到这一点，我正在阅读有关广播的内容，或者可能使用卷积来实现这一点，但我不确定是否有更好的方法来做到这一点。

Answer 1

假设B.shape都是奇数，您可以使用np.indices，操纵它们指向您想要的位置，然后使用np.add。在

def array_sum(A, B, loc = (0, 0)):
    A_ = A.copy()
    ix = np.indices(B.shape)
    new_loc = np.array(loc) - np.array(B.shape) // 2
    new_ix = np.mod(ix + new_loc[:, None, None], 
                    np.array(A.shape)[:, None, None])
    np.add.at(A_, tuple(new_ix), B)
    return A_

测试:

array_sum(A, B)
Out:
array([[ 5.,  6.,  0.,  4.],
       [ 8.,  9.,  0.,  7.],
       [ 0.,  0.,  0.,  0.],
       [ 2.,  3.,  0.,  5.]])