# 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]), &i > 0\\ arr[i, j], &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 > 0:
if sum_arr[i - 1, j - 1] < sum_arr[i - 1, min_index]:
min_index = j - 1
if j < w - 1:
if sum_arr[i - 1, j + 1] < 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 <j - 1 : j + 2) but check if task don't need <j - 1, j + 2>

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

• Storing parent seems unnecessary given we could replace parent[j] with sum_arr[i - 1, j]. Also a (slightly, just about 0.04s for h, w = 1000 in the given code) faster vectorization could be found here: stackoverflow.com/questions/53889590/… – Spellstaker Dec 22 '18 at 6:41