pythonpython-3.xnumpybroadcast

broadcast 2d to 3d with non matching columns


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