python - numpy subtract/add 1d array from 2d array

22
votes

I have the following 2D-array:

a = array([[ 1,  2,  3],
           [ 4,  5,  6],
           [ 7,  8,  9],
           [10, 11, 12],
           [13, 14, 15]])

and another 1D-array:

b = array([ 1,  2,  3,  4,  5])

then I want to calculate something like 

c = a - b

with the intent of getting:

c = array([[0, 1,  2],
           [2, 3,  4],
           [4, 5,  6],
           [6, 7,  8],
           [8, 9, 10]])

but instead I get the error message:

Traceback (most recent call last):
  Python Shell, prompt 79, line 1
ValueError: operands could not be broadcast together with shapes (5,3) (5,)

I read the broadcasting rules but didn´t get any wiser. I could do a workaround with for-loops or similar but there should be a direct way. Thanks

2
pythonarraysnumpy
Do a - b[:,None] or a - b[:,np.newaxis], that is convert b to a 2D array by adding a dimension and then do subtraction. - Divakar
@Divakar you should make an answer out of your comment. - yati sagade
Isn't a basic question like this covered before on SO? Dup-hammer? :) - Divakar

2 Answers

28
votes

You need to convert array b to a (2, 1) shape array, use None or numpy.newaxis in the index tuple. Here is the Indexing of Numpy array.

You can do it Like:

import numpy

a = numpy.array([[ 1,  2,  3],
           [ 4,  5,  6],
           [ 7,  8,  9],
           [10, 11, 12],
           [13, 14, 15]])

b = numpy.array([ 1,  2,  3,  4,  5])
c=a - b[:,None]
print c

Output:

Out[2]: 
array([[ 0,  1,  2],
       [ 2,  3,  4],
       [ 4,  5,  6],
       [ 6,  7,  8],
       [ 8,  9, 10]])
3
votes

As Divakar specified in the comments, just add a new axis to b.

I suggest you read more about broadcasting which is very often useful to vectorize computations in numpy: interestingly enough, a.transpose() - b wouldn't have raised an error (you'd need to transpose the result again to obtain your desired output).

In this computaion, the first array's shape is (3, 5), and b.shape is (5,). So the shape of b corresponds to the tail of the shape of a, and broadcasting can happen. This is not the case when the shape of the first array is (5, 3), hence the error you obtained.

Here are some runtime tests to compare the speeds of the suggested answers, with your values for a and b : you can see that the differences are not really significant

In [9]: %timeit (a.T - b).T
Out[9]: 1000000 loops, best of 3: 1.32 µs per loop

In [10]: %timeit a - b[:,None]
Out[10]: 1000000 loops, best of 3: 1.25 µs per loop

In [11]: %timeit a - b[None].T
Out[11]: 1000000 loops, best of 3: 1.3 µs per loop