查找满足条件的 2D numpy 数组的索引

Question

我有一个大型 2D numpy 数组，并希望找到其中满足条件的 1D 数组的索引：例如，至少有一个大于给定阈值 x 的值。

我已经可以通过以下方式做到这一点，但有没有更短、更有效的方式来做到这一点？

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]

Answer 1

你可以使用 numpy 的内置布尔运算：

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))