CODEWITHSUNDEEP
pythonnumpy
Quasímodo
Quasímodo

Reputation: 4004

Numpy array: efficiently assign values

I have an array and I want to loop through its values to update it as follows:

import numpy as np
arr=np.ones((5,7))
for i in range(1,arr.shape[0]-1):
    for j in range(1,arr.shape[1]-1):
        arr[i,j]=arr[i+1,j]+arr[i,j+1]

The result is, as desired,

[[1. 1. 1. 1. 1. 1. 1.]
 [1. 2. 2. 2. 2. 2. 1.]
 [1. 2. 2. 2. 2. 2. 1.]
 [1. 2. 2. 2. 2. 2. 1.]
 [1. 1. 1. 1. 1. 1. 1.]]

However, for-loops are quite slow and I'd like to know if there is a way to make this more efficient.

Edit: The input is not always np.ones((5,7)), it will be something more heterogeneous in general.

Upvotes: 0

Views: 249

Answers (2)

Kris
Kris

Reputation: 528

Here is the code for question.

import numpy as np 
a=np.random.randn(5, 7)
a1=a
a2=a
mid_mat= a[1:, :][:, :-1]+a[:, :-1][:-1, :]
a1[1:-1, 1:-1] = mid_mat[:-1, :-1]

# Assert Code

for i in range(1,a.shape[0]-1):
    for j in range(1,a.shape[1]-1):
        a2[i,j]=a[i+1,j]+a[i,j+1]

np.testing.assert_array_equal(a1, a2)

Upvotes: 2

jedwards
jedwards

Reputation: 30210

If you draw a box around the "inner" elements, your code is setting the new value of those elements to be the sum of (a) that box "shifted one row down", and (b) that box "shifted one column to the right".

For example:

-----     -----     -----
-XXX-     -----     --XXX
-XXX-  =  -XXX-  +  --XXX
-XXX-     -XXX-     --XXX
-----     -XXX-     -----

And you can do that without loops as follows:

arr[1:-1,1:-1] = arr[2:,1:-1] + arr[1:-1,2:]

Upvotes: 2

Related Questions