broadcast 2d to 3d with non matching columns

March 12, 2024

How can I change the code so it efficiently and without many programming loops subtracts 1D numpy array B from 2D numpy array A in a 3rd dimension So I get C[0, i, j] = A[i,j] – B[0] and 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]])

>Solution :

You need to broadcast B to 3D:

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

Output:

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]]])

Comparison with a loop:

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