Numpy 数组，在特定轴上进行乘法和求和

Question

Let's say I have an array A such that

A = np.array([
    [[1,1,1,1],[2,2,2,2]],
    [[3,3,3,3],[4,4,4,4]],
    [[5,5,5,5],[6,6,6,6]]
])

where the shape is something like (3,2,4) in this example. Now let's say I have another array B such that

B = np.array([1,2,3,4])

I would like to multiply A and B element wise along the last axis of A and sum, such that

C = np.array([
    [10,20],
    [30,40],
    [50,60]
])

Is there a nice way to do this? I thought about making an equivalent 3D array out of B, doing element wise multiplication, and summing along the last axis of this new array. I was wondering if there is a cleaner way to do this?

Edit: If it makes things easier, A can also be written just as

A = np.array([
    [1,2],
    [3,4],
    [5,6]
])

I made A in the way shown at the top of my post because I thought it was neccessary to do so for the proposed multiplication and summation. If working with this version of A at the bottom of this post is easier/just as easy then that would be preferable.

Cheers

Answer 1

Your example is ambiguous and how you break down the computation is unclear.

It looks to me that you could take any element of the last dimension of A and multiply it with the sum of B:

A[...,0]*B.sum()

Or do you want to sum afterwards?

(A*B[None,:]).sum(2)

output:

array([[10, 20],
       [30, 40],
       [50, 60]])

Answer 2

In [126]: A = np.array([
     ...:     [[1,1,1,1],[2,2,2,2]],
     ...:     [[3,3,3,3],[4,4,4,4]],
     ...:     [[5,5,5,5],[6,6,6,6]]
     ...: ])
In [127]: A.shape
Out[127]: (3, 2, 4)
In [128]: B = np.array([1,2,3,4])
 

(3,2,4) broadcasts with (4,) (eg (1,1,4)) to make (3,2,4), then sum on 4:

In [129]: (A*B).sum(axis=-1)
Out[129]: 
array([[10, 20],
       [30, 40],
       [50, 60]])

If the start is (3,2):

In [130]: A1 = np.array([
     ...:     [1,2],
     ...:     [3,4],
     ...:     [5,6]
     ...: ])
     ...: 

make it (3,2,1):

In [132]: (A1[:,:,None]*B).shape
Out[132]: (3, 2, 4)
In [133]: (A1[:,:,None]*B).sum(axis=-1)
Out[133]: 
array([[10, 20],
       [30, 40],
       [50, 60]])