numpy summation of elementwise product along specific direction

user_14513243 0
Viewed 24

I have two different dimension numpy array, lets say A and B, defined as follows

A = np.random.rand(3,3,10)
B = np.random.rand(3,3)

I'm trying to calculate the sum of elementwise product between A and B along the third dimension of A.

ans = []
for i in range(10):
    ans.append(np.sum(B*A[:,:,i]))

Is there any better way to do this? cause I feel its slow when data gets larger.

numpy-ndarray numpy python
1 Answers

Try broadcasting the missing dimension:

(A * B[...,None]).sum(axis=(0,1))
user_42384080
Related