CODEWITHSUNDEEP
pythonnumpy
Wilson
Wilson

Reputation: 666

Find the numpy array elements satisfying the condition

I have a numpy array and I want to find the elements satisfying my condition The codes are as below:

import numpy as np
a = np.array([[1, 2], [1, 3], [1, 2]])
b = np.array([1, 2])
c = (a == b)

The results are

[[ True  True]
[ True False]
[ True  True]]

But what I want is [True, False, True] or the indices [0, 2].

Although I can achieve this by list comprehension, like

c = [all(b==x) for x in a]

But I want to find this element in a 3d matrix in the future, like

a = np.array([[[1, 2], [1, 3], [1, 2]], 
         [[7, 2], [1, 2], [4, 2]]])

I want the index like [[0, 0], [0, 2], [1, 1]]

How should I achieve this by numpy?

Upvotes: 2

Views: 1205

Answers (1)

Stephen Rauch
Stephen Rauch

Reputation: 49784

How about numpy.all with axis:

Code:

np.all(a == b, axis=1)

Test Code:

a = np.array([[1, 2], [1, 3], [1, 2]])
b = np.array([1, 2])
c = np.all(a == b, axis=1)
print(c)

Results:

[ True False  True]

Upvotes: 2

Related Questions