Hi everyone and thank you for your assistance. I am new to python and failed to find an efficient alternative to for loops for the following task.
I want to multiply ndarrays A
and B
of dimension (d,n,m)
and (d,m)
, respectively. With some abuse of terminology to help understanding, A
is a list of nxm
matrices and B
is a list of vectors in R^m
.
For example:
A = np.array([[[0,0,0,0,0],[1,1,1,1,1],[2,2,2,2,2]],[[3,3,3,3,3],[4,4,4,4,4],[5,5,5,5,5]]])
B = np.array([[1,2,3,4,5],[5,6,7,8,9]])
My solution uses a for loop
for i in range(2):
print(A[i]*B[i])
Is there any cheaper alternative (no loops)?
Thank you again
CodePudding user response:
In this case, you can use broadcasting by adding in a new dimension in the "middle" for B
:
>>> import numpy as np
>>> A = np.array([[[0,0,0,0,0],[1,1,1,1,1],[2,2,2,2,2]],[[3,3,3,3,3],[4,4,4,4,4],[5,5,5,5,5]]])
>>> B = np.array([[1,2,3,4,5],[5,6,7,8,9]])
>>> A * B[:, None, :]
array([[[ 0, 0, 0, 0, 0],
[ 1, 2, 3, 4, 5],
[ 2, 4, 6, 8, 10]],
[[15, 18, 21, 24, 27],
[20, 24, 28, 32, 36],
[25, 30, 35, 40, 45]]])
Here is a link to the official docs
Note, your original solution already relied on broadcasting:
>>> A[0]
array([[0, 0, 0, 0, 0],
[1, 1, 1, 1, 1],
[2, 2, 2, 2, 2]])
>>> B[0]
array([1, 2, 3, 4, 5])
>>> A[0] * B[0]
array([[ 0, 0, 0, 0, 0],
[ 1, 2, 3, 4, 5],
[ 2, 4, 6, 8, 10]])