使用不匹配的列将 2d 广播到 3d

Question

我怎样才能改变代码，以便它有效地并且没有很多编程循环从第三维的 2D numpy 数组 A 中减去 1D numpy 数组 B 所以我得到 C[0, i, j] = A[i,j] - B[ 0] 且 C[1,i,j] = A[i,j] - B[1]。

import numpy as np
B=np.array([1, 100])
A=np.arange(4*5).reshape(4,5)
#C=A-B this will not work as A and B have different number of columns
#A=array([[ 0,  1,  2,  3,  4],
#        [ 5,  6,  7,  8,  9],
#        [10, 11, 12, 13, 14],
#        [15, 16, 17, 18, 19]])

Answer 1

您需要将

B

广播为3D：

C = A - B[:,None,None]

输出：

array([[[  -1,    0,    1,    2,    3],
        [   4,    5,    6,    7,    8],
        [   9,   10,   11,   12,   13],
        [  14,   15,   16,   17,   18]],

       [[-100,  -99,  -98,  -97,  -96],
        [ -95,  -94,  -93,  -92,  -91],
        [ -90,  -89,  -88,  -87,  -86],
        [ -85,  -84,  -83,  -82,  -81]]])

与循环比较：

I, J = A.shape

C2 = np.zeros((2, I, J))

for i in range(I):
    for j in range(J):
        C2[0, i, j] = A[i,j] - B[0]
        C2[1,i,j] = A[i,j] - B[1]
        
np.allclose(C, C2)
# True