Apply a 1D array to every cell of a 2D array to create a 3D array

Question

Given a 2D array and a 1D array in Numpy:

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

I'd like to subtract a - b but instead of getting:

Out[16]: 
array([[-1, -2, -3],
   [ 2,  1,  0]])

I want this instead:

array([[[-1, -3, -5], [0, -2, -4], [1, -1, -3]],
       [[2, 0, -2], [3, 1, -1], [4, 2, 0]]])

How do I do that vectorized (i.e. no for loops) for speed?

Answer 1

One solution is

a[...,np.newaxis]-b 

This is just a broadcasting trick.

Instead of subtracting from

>>> a
array([[1, 2, 3],
       [4, 5, 6]])

you subtact from

>>> a[...,np.newaxis]
array([[[1],
        [2],
        [3]],

       [[4],
        [5],
        [6]]])

which has shape

>>> a[...,np.newaxis].shape
(2, 3, 1)

while b has shape

>>> b.shape
(3,)

So broadcasting rules will give you an output with shape (2,3,3) like you want. There is good information on the principles of broadcasting here.