返回2D numpy数组中满足条件的数组的索引

时间:2019-05-21 14:38:02

标签: python arrays numpy matrix

我有一个大型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]

1 个答案:

答案 0 :(得分:4)

您可以使用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))