Search code examples
pythonarraysnumpyboolean-operations

Numpy boolean indexing if number is in list

I have the following array:

x = np.array([
    [2, 0],
    [5, 0],
    [1, 0],
    [8, 0],
    [6, 0]])

I've learned that you can use boolean operations to change selected values in a numpy array. If I want to change the value of the 2nd column to 1 for the rows where the 1st value is equal to 2, 5 or 8 I can do the following:

x[x[:, 0] == 2, 1] = 1
x[x[:, 0] == 5, 1] = 1
x[x[:, 0] == 8, 1] = 1

Which changes the output to:

[[2 1]
 [5 1]
 [1 0]
 [8 1]
 [6 0]]

If that were "normal" python code, I know I could do:

if value in [2, 5, 8]: ...

Instead of:

if value == 2 or value == 5 or value == 8: ...

Is there a shorthand to do something like this with numpy arrays?

Solution

  • You can use numpy's isin method:

    x[np.isin(x[:, 0], [2, 5, 8]), 1] = 1