Numpy array modifying multiple elements at once

Question

I have three numpy arrays:

row = np.array([1,2,3,4,5])

# a is a subset of row:

a = np.array([1, 5])

# b is an array that I use to change some elements in the first row array:

b = np.array([10, 550])

What I need to do is to change in one shot the elements of the row array that are present in a with the correspondent b elements.

i.e.:

>> modified_row
array([10, 2, 3, 4, 500])

Doing this in a naive way would be:

for i in range(len(a)):
    row[np.where(row==a[i])]= b[i]

I would like a solution like;

row[np.where(row==a)] = b

But that doesn't work...

Thanks in advance!

Answer 1

If you don't have guarantees on the sorting of your arrays, you could have a reasonably efficient implementation using np.searchsorted:

def find_and_replace(array, find, replace):
    sort_idx = np.argsort(array)
    where_ = np.take(sort_idx, 
                     np.searchsorted(array, find, sorter=sort_idx))
    if not np.all(array[where_] == find):
        raise ValueError('All items in find must be in array')
    row[where_] = b

The only thing that this can't handle is repeated entries in array, but other than that it works like a charm:

>>> row = np.array([5,4,3,2,1])
>>> a = np.array([5, 1])
>>> b = np.array([10, 550])
>>> find_and_replace(row, a, b)
>>> row
array([ 10,   4,   3,   2, 550])

>>> row = np.array([5,4,3,2,1])
>>> a = np.array([1, 5])
>>> b = np.array([10, 550])
>>> find_and_replace(row, a, b)
>>> row
array([550,   4,   3,   2,  10])

>>> row = np.array([4, 5, 1, 3, 2])
>>> find_and_replace(row, a, b)
>>> row
array([  4, 550,  10,   3,   2])