CODEWITHSUNDEEP
arraysnumpy
TheChymera
TheChymera

Reputation: 17934

Assign 1 and 0 values to numpy array depending on whether values are in list

I am looking for a way to filter numpy arrays based on a list

input_array = [[0,4,6],[2,1,1],[6,6,9]]
list=[9,4]
...
output_array = [[0,1,0],[0,0,0],[0,0,1]]

I am currently flattening the array, and turning it to a list and back. Looks very unpythonic:

    list=[9,4]
    shape = input_array.shape
    input_array = input_array.flatten()
    output_array = np.array([int(i in list) for i in input_array])
    output_array = output_array.reshape(shape)

Upvotes: 2

Views: 52

Answers (1)

Divakar
Divakar

Reputation: 221624

We could use np.in1d to get the mask of matches. Now, np.in1d flattens the input to 1D before processing. So, the output from it is to be reshaped back to 2D and then converted to int for an output with 0s and 1s.

Thus, the implementation would be -

np.in1d(input_array, list).reshape(input_array.shape).astype(int)

Sample run -

In [40]: input_array
Out[40]: 
array([[0, 4, 6],
       [2, 1, 1],
       [6, 6, 9]])

In [41]: list=[9,4]

In [42]: np.in1d(input_array, list).reshape(input_array.shape).astype(int)
Out[42]: 
array([[0, 1, 0],
       [0, 0, 0],
       [0, 0, 1]])

Upvotes: 2

Related Questions