In numpy, how to compare all values in an axis

Viewed 3293

For a numpy array, how can I change the value only if all elements along an axis are equal to another array? For example...

array = np.array([[1, 0, 1],
                  [0, 0, 1],
                  [1, 1, 0],
                  [0, 0, 0],
                  [1, 0, 1]])

I want to replace all [1, 0, 1] with [1, 1, 1]... so that array becomes

array([[1, 1, 1],
       [0, 0, 1],
       [1, 1, 0],
       [0, 0, 0],
       [1, 1, 1]])

When I use a boolean array, it checks each individual number. How can I compare the entire row at once instead?

4 Answers

Try with:

array[(array == [1, 0, 1]).all(axis=1)] = [1, 1, 1]

Your array, and search array:

In [87]: arr = np.array([[1, 0, 1], 
    ...:                   [0, 0, 1], 
    ...:                   [1, 1, 0], 
    ...:                   [0, 0, 0], 
    ...:                   [1, 0, 1]])                                                  
In [88]: x = np.array([1,0,1])                                                          

WIth a (5,3) and (3,) shape we can compare this with:

In [89]: arr == x                                                                       
Out[89]: 
array([[ True,  True,  True],
       [False,  True,  True],
       [ True, False, False],
       [False,  True, False],
       [ True,  True,  True]])

This makes use of broadcasting to treat the x as (1,3). In other cases we may need to adjust the arrays for proper broadcasting.

And we can find the rows that match exactly, all True:

In [90]: (arr == x).all(axis=1)                                                         
Out[90]: array([ True, False, False, False,  True])

If we use that boolean array to index the first dimension of arr, we get the [1,0,1] rows:

In [96]: arr[(arr == x).all(axis=1)]                                                    
Out[96]: 
array([[1, 0, 1],
       [1, 0, 1]])

and we can use the same indexing to replace the values of those rows (again this uses broadcasting to put a (3,) array into a (2,3) slot:

In [97]: arr[(arr == x).all(axis=1)]=[1,1,1]                                            
In [98]: arr                                                                            
Out[98]: 
array([[1, 1, 1],
       [0, 0, 1],
       [1, 1, 0],
       [0, 0, 0],
       [1, 1, 1]])

This is a good exercise in numpy broadcasting, indexing, and axis reduction.

Related