Numpy: How can I add (4, ) array and (9, 9) array to (9, 9, 4) array?

user_10516199 0
Viewed 76
import numpy as np
a = np.array([0, 10, 11, 1])
 
# OUT:
# array([ 0, 10, 11,  1])
 
b = np.arange(81).reshape(9, 9)
# OUT:
# array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8],
#       [ 9, 10, 11, 12, 13, 14, 15, 16, 17],
#       [18, 19, 20, 21, 22, 23, 24, 25, 26],
#       [27, 28, 29, 30, 31, 32, 33, 34, 35],
#       [36, 37, 38, 39, 40, 41, 42, 43, 44],
#       [45, 46, 47, 48, 49, 50, 51, 52, 53],
#       [54, 55, 56, 57, 58, 59, 60, 61, 62],
#       [63, 64, 65, 66, 67, 68, 69, 70, 71],
#       [72, 73, 74, 75, 76, 77, 78, 79, 80]])

I want each item of b-array add to a-array, and get the result like, just like a simple loop can do:

result = list()

for row in b:
    ls = list()
    
    for i in row:
        c = a + i
        ls.append(c)

    result.append(ls)
np.array(result)

Loop is not efficient enough, is there any way that I can use the Numpy broadcast rule to get this result?(Numpy array)

pandas matrix numpy arrays python
1 Answers

The kind of "general" way to do this is by indexing with None. If you index an array with None, it creates a new axis and broadcasts the array along that axis.

So if you have a.shape = (4,) and b.shape = (9,9), you can add like this:

a[:,None,None] + b[None,:,:]

The : in the index for a is the only axis of a, which has size 4. The remaining two axes, the None,None, are new axes that get broadcast to match b.

The :,: indexes in the index for b are the axes in the array, with shape (9,9). The None is a new axis, broadcast to match a.

There may be a shorter way to write this but I can't be bothered to remember how the rules for implicit broadcasting works, so I just make it explicit.

You can also use numpy.newaxis if that is clearer, but it has the same effect.

user_822940
Related