Python/Numpy/Boolean 索引：对于数组中的每个 True 值，修改接下来的 2 个元素

如何解决Python/Numpy/Boolean 索引：对于数组中的每个 True 值，修改接下来的 2 个元素

我有一个 numpy 数组

array = np.array([5,100,5,5])

我创建了一个带有布尔索引的掩码，如下所示：

mask = (array < 30)

这给出了一个面具

[ True False False False  True  True False False False  True]

我可以使用

获取掩码中 True 值的索引

indices = np.where(mask)[0]

这给了

[0 4 5 9]

对于掩码中的每个 True 值，我想将接下来的 2 个元素也修改为 True。

我可以像这样使用 for 循环来做到这一点：

for i in indices:
    mask[i:i+3] = True

在不使用 for 循环的情况下，是否有更 numpythonic 的方法？

所需的掩码输出：

[ True  True  True False  True  True  True  True False  True]

这里的主要优先事项是性能。

解决方法

您可以使用 np.flatnonzero 来简化索引的获取。然后你可以为每一个添加 np.arange(3)：

ind = np.flatnonzero(mask)[:,None] + np.arange(3)

唯一需要注意的是，您的索引可能包含几个越界元素。您可以使用蒙版或 np.clip:

修剪它们

ind[ind >= mask.size] = mask.size - 1

然后您可以直接应用索引，因为 numpy 允许花式索引的任意维度：

mask[ind] = True

如果有少量涂抹，可以直接涂抹面膜：

mask[1:] |= mask[:-1]
mask[1:] |= mask[:-1]

如果涂抹量是任意的，您显然必须将其置于循环中，但您可以通过步进 2 的幂来优化它。

我称此操作为 mask[1:] |= mask[:-1] 涂抹，因为它会将任意一组 True 元素的大小向右扩大一个，就像您用手指涂抹墨水一样。涂抹任意数量的n：

s = 1
while s <= n:
    mask[s:] |= mask[:-s]
    s *= 2
s = n - s // 2
if s:
    mask[s:] |= mask[:-s]