CODEWITHSUNDEEP
pythonnumpy
Shady45
Shady45

Reputation: 45

Conditioning on NumPy array based on values from different arrays

Suppose I have two (sample) arrays in the following manner:

a = np.array([4, 5,-1, 2, -3, 3, -4])
b = np.array([0, 1, 0, 0, 1, 1, 0])

Now, I want to calculate the count of occurrences where (a > 0 and b == 0) and also where (a < 0 and b == 1). How can I condition an array based on values of two different arrays? If I try this:

a[a < 0 and b == 1]

I get the error:

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

I can do it using loops, but I want to avoid it and see if there's a better way of doing this.

Upvotes: 1

Views: 339

Answers (1)

jkr
jkr

Reputation: 19250

One can use the bitwise-and operator &. Take care to wrap the expressions in parentheses because & binds more tightly

a[(a < 0) & (b == 1)]

One can achieve the same behavior with numpy.logical_and.

mask = np.logical_and(a < 0, b == 1)
a[mask]

Upvotes: 4

Related Questions