在 numpy 中仅屏蔽特定范围的索引

Question

我有一个 ndarray。我需要屏蔽任何小于 50 的数字，直到遇到的第一个数字大于 50。这应该在一行的开头和结尾都完成。对，当遇到的第一个数字大于 50 时，屏蔽应该停止。

一行看起来像：

    [ 0   1   1   0  57  121  120  157  77  14   0   3   0   0   0   0  67 100  
    98  97 101 129 139 105  97 105 181 126  10   0   0]

我想要类似的东西：

    [-- -- -- -- 57 121 120 157 77 14  0  3  0  0  0  0 67 100 98 97
 101 129 139 105 97 97 105 181 126 -- -- --]

屏蔽应该在第二行的 57 之前停止，在最后一行的 126 处停止。

我尝试过 ma.masked_where，但它也掩盖了两者之间的 0，这是我不想要的。

所以，如果有办法做到这一点，或者你能帮助我，以便我可以指定一个索引范围，例如：[0:40] 只应该被屏蔽。

我不想在屏蔽后更改数组的维度。此外，-- 的存在不会影响我的目的。

Answer 1

您可以使用布尔索引或手动迭代。前者对小数组效率更高；后者用于具有少量超出范围值的大型数组。

布尔索引

x = np.array([0, 0, 0, 2, 3, 51, 34, 1, 23, 32, 32, 52, 0, 0, 0])

start = (x > 50).argmax()
end = len(x) - (x[::-1] > 50).argmax()

print(x[start: end])

[51 34  1 23 32 32 52]

手动迭代

将next与生成器表达式和enumerate一起使用：

start = next(idx for idx, val in enumerate(x) if val > 50)
end = len(x) - next(idx for idx, val in enumerate(reversed(x)) if val > 50)

print(x[start: end])

[51 34  1 23 32 32 12]

屏蔽

如果你想用np.nan替换超出范围的值，你可以进行相应的赋值，记得先转换为float，因为NaN的值是float：

x = x.astype(float)
x[:start] = np.nan
x[end:] = np.nan

print(x)

array([nan, nan, nan, nan, nan, 51., 34.,  1., 23., 32., 32., 52., nan, nan, nan])