Summing over 2D and 3D arrays using broadcasting

Question

Consider the following MWE:

import numpy as np
n=2
N = 6
a = np.random.randint(0,10,size=(N,1))
b = np.random.randint(0,10,size=(N,n))
c = np.random.randint(0,10,size=(n,N,5))

where c is e.g. (it is random recall):

array([[[7 5 1 7 0]
  [2 8 2 1 4]
  [0 4 1 7 3]
  [1 6 6 9 6]
  [9 6 0 0 2]
  [9 6 0 6 7]]

 [[0 3 9 0 3]
  [4 7 5 3 8]
  [8 0 6 7 9]
  [5 4 9 5 2]
  [5 6 6 8 7]
  [7 7 2 6 0]]])

and has shape (2,6,5).

From which we make:

out = a+b
>>>out
array([[ 9,  7],
       [ 5,  7],
       [ 7,  3],
       [ 9,  9],
       [15, 10],
       [ 8,  9]])

which has shape (6,2).

Now here is what I want to do: I want to add the first column of out to the first matrix of c (i.e. where matrices are indexed by the first dimension of c), the second column of out to the second column of c and so on (you get the drift). Currently, I am attempting to do this using broadcasting but I seem to have confused myself.

I want to do without using loops as my real problem is very large.

Desired output:

>>>np.stack([out[:,i][:,np.newaxis] + c[i] for i in range(2)])

array([[[16, 14, 10, 16,  9],
        [ 7, 13,  7,  6,  9],
        [ 7, 11,  8, 14, 10],
        [10, 15, 15, 18, 15],
        [24, 21, 15, 15, 17],
        [17, 14,  8, 14, 15]],

       [[ 7, 10, 16,  7, 10],
        [11, 14, 12, 10, 15],
        [11,  3,  9, 10, 12],
        [14, 13, 18, 14, 11],
        [15, 16, 16, 18, 17],
        [16, 16, 11, 15,  9]]])

which has shape (2,6,5).

Attempt:

out[None, :,:] + c

which renders the following error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-702-d54cfe51ec32> in <module>
----> 1 out[None, :,:] + c

ValueError: operands could not be broadcast together with shapes (1,6,2) (2,6,5)

Help would be most appreciated.

Answer 1

You can transpose and add a dimension and let the broadcasting do the job:

out.T[...,None]+c 

Explanation: .T transposes out (to shape (2,6)) and [...,None] adds an extra dimension as the latest dimension of out (Now out is of shape (2,6,1)). Finally, the broadcasting to c with shape (2,6,5) will broadcast all elements to the depths of c as desired.