一个掩码数组索引问题

Question

我有一个包含一些NaN值的NumPy数组：

arr = [ 0, NaN, 2, NaN, NaN, 5, 6, 7 ]

使用一些逻辑（超出问题范围），我生成了一个NaN位置的掩码:

mask = [ True, False, True, False, False, True, True, True ]

我使用这个掩码来选择有效数据：

valid_arr = arr[mask]   # [ 0, 2, 5, 6, 7 ]

然后我执行一个任意的算法，选择了这个新数组中的几个 索引值：

indeces = myAlgo(valid_arr)   # [ 1, 3 ]

有效数组中的索引为1,3（对应值为2和6）。我需要知道这些索引在原始数组（arr）中对应的是哪个索引。在上面的例子中，显然是2和6。

该数组是时间序列数据，未排序。一种解决方法是遍历掩码，仅在找到有效数字时递增计数器。是否可以使用numpy更有效地完成此操作？

- David Ferris

Answer 1

你可以使用返回原始数组索引的掩码来压平它，然后使用新的索引来获取掩码索引的子集：

mask = np.array([ True, False, True, False, False, True, True, True ])
indices = [1,3]

np.flatnonzero(mask)[indices]
# array([2, 6])

- Psidom

---

我对这个如何工作有点困惑。根据numpy文档，flatnonzero(a)返回在a的扁平化版本中非零的索引。我认为这也适用于True/False掩码数组？ - David Ferris

是的。在这里，True被视为1，而False被视为0，就像大多数隐式类型转换发生的情况一样。 - Psidom

啊，我明白了...它包含与原始数组对应的索引！太棒了！ - David Ferris

Answer 2

您可以创建一个索引数组（使用np.arange），然后对其进行掩码处理，并最终使用整数数组高级索引来获取相应的项：

>>> mask = np.array([ True, False, True, False, False, True, True, True ])
>>> np.arange(mask.size)[mask][[1, 3]]
array([2, 6])

- MSeifert