Cumulative sum of the smallest neighbors in the previous row in a Numpy array I have the following code snippet, which essentially does the following:
Given a 2d numpy array, arr, compute sum_arr  as follow:
$$sum\_arr[i, j] = \begin{cases}
    arr[i, j] + min(sum\_arr[i - 1, j-1:j+2]), &amp;i &gt; 0\\ arr[i, j], &amp;i = 0\end{cases}$$ 
(reasonable indices for j - 1 : j + 2 of course, all within 0 and w)  

Here's my implementation:

import numpy as np

h, w = 1000, 1000 # Shape of the 2d array
arr = np.arange(h * w).reshape((h, w)) 

sum_arr = arr.copy()

def min_parent(i, j):
    min_index = j    
    if j &gt; 0:
        if sum_arr[i - 1, j - 1] &lt; sum_arr[i - 1, min_index]:
            min_index = j - 1
    if j &lt; w - 1:
        if sum_arr[i - 1, j + 1] &lt; sum_arr[i - 1, min_index]:
            min_index = j + 1    
    return (i - 1, min_index)


for i, j in np.ndindex((h - 1, w)):
    sum_arr[i + 1, j] += sum_arr[min_parent(i + 1, j)]


And here's the problem: this code snippet takes way too long to execute for only 1e6 operations (About 5s on average on my machine)

What is a better way of implementing this?
first, your problem is example of dynamic programming. Here some tips how to approach problem in python
your solution calculate each table cell separately. But function min_parent() use data only from previous row, so you can calculate row by row (and numpy is designed to fast operations on vectors)
i am using like you range &lt;j - 1 : j + 2) but check if task don't need &lt;j - 1, j + 2&gt;


Example code


assert(w >= 2) # if not, you should handle w == 1 in special way

def vectorized_solution(arr, h):
    sum_arr = arr.copy()
    for i in range(1, h):
        parent = sum_arr[i - 1, :]
        sum_arr[i, 0] += min(parent[0], parent[1])
        sum_arr[i, 1:-1] += np.minimum.reduce([parent[0:-2], parent[1:-1], parent[2:]])
        sum_arr[i, -1] += min(parent[-2], parent[-1])
    return sum_arr