How to set different values of the elements of a np.arrays to different values in Python 3.8?-CodePudding

Let I have the following np.array:

>>>a=np.array([20, 10,5,10,5,10])
>>>array([20, 10,  5, 10,  5, 10])

Now, I want to replace 20 and 10 by 1 and 5 by 0.

Is there a function that can do that in one step?

Here is what I have tried:

>>>a[a==10]=1
>>>a[a==10]=1
>>>a[a==5]=0

and I am getting my desired output, which is:

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

As you can see, I had to follow three steps in order to get my result. But I want to get my result only in one step. Is there a function that can deliver my result in one step?

CodePudding user response：

You can use the map function.

list(map(lambda x: int(x in [10,20]),a))

The map function will apply the function in the first argument to all the elements in the list given as the second argument. Here the lambda function returns 0 if the element is not 10 or 20, and 1 if the element is 10 or 20.

EDIT FOLLOWING THE AUTHOR'S COMMENT

To keep the result as a numpy array, you can use the from iter numpy function :

a = np.fromiter(map(lambda x: int(x in [10,20]),a),dtype=int)

CodePudding user response：

Despite @robinood's answer works fine, I'd prefer this way:

import pandas as pd
a = pd.Series(a).replace([20,10,5],[1,1,0]).values

This because, for long arrays, cycling can take a great amount of time. For this reason i tested both mine and @robinood's solution on a = np.random.choice([20, 10, 5], size=10_000_000) and the results are the following:

my solution: 655 ms ± 7.09 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
@robinood's solution 9.51 s ± 25.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)