在numpy数组中查找小于x的元素的索引

Question

假设我有一个 numpy 数组，例如：

import numpy as np    
arr = np.array([10,1,2,5,6,2,3,8])

如何提取包含小于 6 的元素索引的数组，以便得到以下结果：

np.array([1,2,3,5,6])

我想要一些行为类似于np.nonzero() 的东西，但不是测试非零值，而是测试小于 x 的值

Answer 1

您可以在 a 的扁平版本中使用非零numpy.flatnonzero的布尔掩码和返回索引：

np.flatnonzero(arr < 6)
# array([1, 2, 3, 5, 6])

一维数组的另一个选项是numpy.where：

np.where(arr < 6)[0]
# array([1, 2, 3, 5, 6])