Using an index array to assign values to an array in numpy

Question

I am trying to vectorize this operation to speed up run time. To set up my problem look at the following code.

current_array=np.array([[2,3],[5,6],[34,0]])
ind =np.array([[0,1],[1,1],[0,0]])
new_vals=[99,77]

##do some magic
result = [[99,77],[77,77],[99,99]]

So we have the current array and I want to use the values in ind to assign the values of new_vals to current_array such that you end up with result.

I have a way to do this but it uses two loops and I would like to speed it up as much as possible. Right now I am doing the following

def set_image_vals(image,means,mins):
    for i in range(0,image.shape[0]):
        for j in range(0,image.shape[1]):
            image[i,j]=means[mins[i,j]]
    return image

where image would be current_array , means would be new_vals and mins would be ind.

Is there a better way to do this that can speed things up?

Answer 1

For the general case, this is the most flexible and fastest:

>>> new_vals = np.array([99, 77])
>>> new_vals[ind]
array([[99, 77],
   [77, 77],
   [99, 99]])

Here, new_vals could have more than two elements, and ind can index up to that number of elements.