7

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]
Reveille
  • 4,359
  • 3
  • 23
  • 46
  • simpler index retrieval cases [1](https://stackoverflow.com/questions/27175400/how-to-find-the-index-of-a-value-in-2d-array-in-python), [2](https://stackoverflow.com/questions/432112/is-there-a-numpy-function-to-return-the-first-index-of-something-in-an-array), [3](https://stackoverflow.com/questions/18079029/index-of-element-in-numpy-array) – Reveille Aug 19 '19 at 15:52

1 Answers1

8

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))
Chris Mueller
  • 6,490
  • 5
  • 29
  • 35