8
\$\begingroup\$

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?

\$\endgroup\$
4
\$\begingroup\$
  • 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
\$\endgroup\$
  • \$\begingroup\$ 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/… \$\endgroup\$ – Spellstaker Dec 22 '18 at 6:41

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.