Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Find indices of 2D numpy arrays that meet a condition

I have a large 2D numpy array and want to find the indices of the 1D arrays inside it that meet a condition: e.g., have at least a value greater than a given threshold x.

I already can do it the following way but is there a shorter, more efficient way to do it?

import numpy

a = numpy.array([[1,2,3,4,5], [1,2,3,4,20], [1,2,2,4,5]])

indices = []
i = 0
x = 10
for item in a:
    if any(j > x for j in item):
        indices.append(i)
    i += 1

print(indices) # gives [1]
like image 240
Reveille Avatar asked May 21 '19 14:05

Reveille


1 Answers

You could use numpy's built-in boolean operations:

import numpy as np
a = np.array([[1,2,3,4,5], [1,2,3,4,20], [1,2,2,4,5]])

indices = np.argwhere(np.any(a > 10, axis=1))
like image 58
Chris Mueller Avatar answered Oct 02 '22 12:10

Chris Mueller