Let’s say I have a NumPy array a with shape (3, 5, 7) and a 1D NumPy array l with shape (3, 1). I want to add the 0th element of l to the 0th subarray in a, the 1st element of l to the 1st subarray of a, and the 2nd element of l to the 2nd and last subarray of a.
For instance, if l = [[30 40 50]], then I want a[0] += 30, a[1] += 40, a[2] += 50.
I know that to achieve this, I can easily do the following:
for i in range(a.shape[0]):
a[i] += l[i]
but I was wondering if there’s a way to do this with pure NumPy, no for loops required. Trying np.add(a, l) invariably creates the error operands could not be broadcast together with shapes (3,5,7) (3,1) and trying the following:
def add_arrs(a, l):
return a + l
a = np.apply_along_axis(add_arrs, 0, a, l)
gives me a final shape of (3, 3, 5, 7), which is not what I want. Thanks for your help!
EDIT: I figured it out:
a = np.apply_along_axis(np.add, 0, a, l.T)[0]
Advertisement
Answer
In [7]: a = np.arange(3*5*7).reshape(3,5,7) In [8]: b = np.array([[30,40,50]]) In [9]: b.shape Out[9]: (1, 3)
That leading 1 dimension doesn’t help; you have to transpose anyways
In [10]: b = np.array([30,40,50]) In [11]: b.shape Out[11]: (3,)
The broadcasting works fine if b is (3,1,1) shape:
In [12]: a += b[:,None,None]
In [13]: a
Out[13]:
array([[[ 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]],
[[ 75, 76, 77, 78, 79, 80, 81],
[ 82, 83, 84, 85, 86, 87, 88],
[ 89, 90, 91, 92, 93, 94, 95],
[ 96, 97, 98, 99, 100, 101, 102],
[103, 104, 105, 106, 107, 108, 109]],
[[120, 121, 122, 123, 124, 125, 126],
[127, 128, 129, 130, 131, 132, 133],
[134, 135, 136, 137, 138, 139, 140],
[141, 142, 143, 144, 145, 146, 147],
[148, 149, 150, 151, 152, 153, 154]]])