python - 在 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

您可以在 bool 掩码上使用 numpy.flatnonzero 并在 a 的扁平化版本中返回非零索引:

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

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

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