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!

Solution 1

Please note @Jaime's answer is better - this was marked as correct despite it relying on the ordering of the elements. Here is a working version that does not modify row in-place but otherwise will work in the general case. At the end of this post is my original answer.

import numpy as np

def replaced(row, a, b):
    row_order = row.argsort()
    a_order = a.argsort()

    sorted_row = row[row_order]
    sorted_a = a[a_order]
    sorted_b = b[a_order]

    sorted_row[np.in1d(sorted_row, sorted_a)] = sorted_b

    # return results in original order
    return sorted_row[row_order.argsort()]

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

row = np.array([1,2,3,4,5])
print replaced(row, a, b)

row = np.array([5,4,3,2,1])
print replaced(row, a, b)

row = np.array([4, 5, 1, 3, 2])
print replaced(row, a, b)

results:

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

ORIGINAL INCORRECT ANSWER

One way to do this is with the in1d function, which will generate a boolean array that you can use to index row as shown below.

Note that you may have problems with this (and other methods) if the elements of row are not unique or if you have repeated elements in a

>>> import numpy as np
>>> row = np.array([1,2,3,4,5])
>>> a = np.array([1, 5])
>>> b = np.array([10, 550])
>>> np.in1d(row, a)
array([ True, False, False, False,  True], dtype=bool)
>>> row[np.in1d(row, a)] = b
>>> row
array([ 10,   2,   3,   4, 550])

You can normally use whatever index/boolean array you originally used to extract a for this purpose too.

Solution 2

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])

Solution 3

Another possibility:

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

The way this works is:

>>> row.reshape(-1, 1) == a
array([[ True, False],
       [False, False],
       [False, False],
       [False, False],
       [False,  True]], dtype=bool)
>>> np.any(row.reshape(-1, 1) == a, axis=1)
array([ True, False, False, False,  True], dtype=bool)

And this boolean mask corresponds to the entries you want to replace.

The time and space complexity of this solution are pretty bad: Θ(nm) to replace m entries in an array of size n due to the large boolean mask. I don't recommend it over in1d for your specific use case, but it shows a detour that is useful in related cases.

Solution 4

An interesting alternative solution is to use numpy.put as documented here. In this case it is also important to think carefully about what will happen if there are duplicates in row. By default, put will cycle through the elements in b if there are more than two matches in this case.

import numpy as np
row = np.array([1,2,3,4,5])
a = np.array([1, 5])
b = np.array([10, 550])
index_list = [np.where(row == element) for element in a]
np.put(row,index_list,b)
row
array([ 10,   2,   3,   4, 550]) #output

Edit: additional example to deal with index-based assignment query in comments:

>>> import numpy as np
>>> target_array = np.arange(50)
>>> n = 2 
>>> index_array = np.arange(0,len(target_array),n)
>>> b = np.array([10, 550])
>>> np.put(target_array, index_array, b)
>>> target_array #note that np.put cycles through the substitutions in b
array([ 10,   1, 550,   3,  10,   5, 550,   7,  10,   9, 550,  11,  10,
        13, 550,  15,  10,  17, 550,  19,  10,  21, 550,  23,  10,  25,
       550,  27,  10,  29, 550,  31,  10,  33, 550,  35,  10,  37, 550,
        39,  10,  41, 550,  43,  10,  45, 550,  47,  10,  49])

Solution 5

You can now use array.put

>>> a = np.arange(5)
>>> np.put(a, [0, 2], [-44, -55])
>>> a
array([-44,   1, -55,   3,   4])

View more solutions