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