Replacing Numpy elements if condition is met Replacing Numpy elements if condition is met python python

Replacing Numpy elements if condition is met


>>> import numpy as np>>> a = np.random.randint(0, 5, size=(5, 4))>>> aarray([[4, 2, 1, 1],       [3, 0, 1, 2],       [2, 0, 1, 1],       [4, 0, 2, 3],       [0, 0, 0, 2]])>>> b = a < 3>>> barray([[False,  True,  True,  True],       [False,  True,  True,  True],       [ True,  True,  True,  True],       [False,  True,  True, False],       [ True,  True,  True,  True]], dtype=bool)>>> >>> c = b.astype(int)>>> carray([[0, 1, 1, 1],       [0, 1, 1, 1],       [1, 1, 1, 1],       [0, 1, 1, 0],       [1, 1, 1, 1]])

You can shorten this with:

>>> c = (a < 3).astype(int)


>>> a = np.random.randint(0, 5, size=(5, 4))>>> aarray([[0, 3, 3, 2],       [4, 1, 1, 2],       [3, 4, 2, 4],       [2, 4, 3, 0],       [1, 2, 3, 4]])>>> >>> a[a > 3] = -101>>> aarray([[   0,    3,    3,    2],       [-101,    1,    1,    2],       [   3, -101,    2, -101],       [   2, -101,    3,    0],       [   1,    2,    3, -101]])>>>

See, eg, Indexing with boolean arrays.


The quickest (and most flexible) way is to use np.where, which chooses between two arrays according to a mask(array of true and false values):

import numpy as npa = np.random.randint(0, 5, size=(5, 4))b = np.where(a<3,0,1)print('a:',a)print()print('b:',b)

which will produce:

a: [[1 4 0 1] [1 3 2 4] [1 0 2 1] [3 1 0 0] [1 4 0 1]]b: [[0 1 0 0] [0 1 0 1] [0 0 0 0] [1 0 0 0] [0 1 0 0]]