AMR1015
AMR1015

Reputation: 5

How do I extract rows of a 2D NumPy array by condition?

I have a 4*5 NumPy array and I want to retrieve rows in which all elements are less than 5.

arr = np.array([[0,2,3,4,5],[1,2,4,1,3], [2,2,5,4,6], [0,2,3,4,3]])
arr[np.where(arr[:,:] <= 4)] 

expected output:

[[1,2,4,1,3],[0,2,3,4,3]]

actual output:

array([0, 2, 3, 4, 1, 2, 4, 1, 3, 2, 2, 4, 0, 2, 3, 4, 3])

Any help is appreciated!

Upvotes: 0

Views: 415

Answers (1)

user17242583
user17242583

Reputation:

This actually quite simple. Just convert the entire array to booleans (each value is True if it's less than 5, False otherwise), and use np.all with axis=1 to return True for each row where all items are True:

>>> arr[np.all(arr < 5, axis=1)]
array([[1, 2, 4, 1, 3],
       [0, 2, 3, 4, 3]])

Upvotes: 1

Related Questions