Numpy select elements with a condition along axis

Viewed 1067

I have a 2D numpy array x as:

[  [ 1,  2,  3],
   [ 4,  5,  6],
   [ 7,  8,  9],
   [10, 11, 12],
   [13, 14, 15],
   [16, 17, 18],
   [19, 20, 21],
   [22, 23, 24],
   [25, 26, 27],
   [28, 29, 30],
   [31, 32, 33],
   [34, 35, 36],
   [37, 38, 39],
   [40, 41, 42],
   [43, 44, 45],
   [46, 47, 48],
   [49, 50, 51],
   [52, 53, 54],
   [55, 56, 57],
   [58, 59, 60]  ]

I want to extract the arguments of rows for which any element in the row is less than 25. So, what I need for output is [0,1,2,3,4,5,6,7] for just the rows but using np.where(x<35) is giving me the list of 2D arguments for all possible values. In other words, I what I want are the arguments of all the rows of x where at least one element is less than 25, but what I am getting are the arguments of all the elements of x that are less than 25. What should I do? Is there a specific function for this or should I just select the unique values from the returned list of arguments?

1 Answers

One way would be this:

import numpy as np
# x is your array
x1 = (x < 25).sum(axis = 1)
rows = np.where(x1 > 0)[0]

The row indices are in rows as array([0, 1, 2, 3, 4, 5, 6, 7]).

You can also use nonzero as:

rows = np.nonzero((x < 25).sum(axis = 1))[0]
Related