Reputation: 5191
Can I create a new array by comparing arrays with numpy?
I have 3 arrays (A1, A2, A3
).
How can I find all indexes where A1 == 2 and A2 > A3
and write there the value 5
in a new array?
I have this matlab code which is doing this:
index = find(A1==2 & A2>A3);
new_array(index) = 5;
I found putmask and logical_and but not sure if this are the right tools and how do use it in my case. Thanks!
Upvotes: 0
Views: 69
Reputation: 11
You can use the np.where function
import numpy as np
A1 = np.array([0, 0, 2, 2, 4, 4])
A2 = np.arange(len(A1))
A3 = np.ones(len(A1))*3
out = np.where((A1 == 2) & (A2 >= A3), 5, 0)
or more simple
out = ((A1 == 2) & (A2 >= A3))*5
Upvotes: 1
Reputation: 53698
The code below uses &
to chain together conditionals. As such whenever A1 == 2
and A2 > A3
are both True
, the index
array will be True
import numpy as np
A1 = np.array([0, 0, 2, 2, 4, 4])
A2 = np.arange(len(A1))
A3 = np.ones(len(A1))*3
new_array = np.zeros(len(A1))
index = (A1 == 2) & (A2 > A3)
new_array[index] = 5
# array([ 0., 0., 5., 5., 0., 0.])
You could use np.logical_and
of course. But this restricts you to only two conditionals, whilst you can effectively chain as many as you like when using &
.
Upvotes: 2