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

Question

我有一个 numpy 数组

array = np.array([5,100,100,100,5,5,100,100,100,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]

这里的首要任务是性能。

Answer 1

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

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

唯一需要注意的是您的索引可能包含几个越界元素。你可以 trim 他们带着面具或 np.clip:

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

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

mask[ind] = True

如果要做小涂抹，可以直接涂抹面膜：

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

如果涂抹量是任意的，你显然必须把它放在一个循环中，但你可以通过以二的幂步进来优化它。

我称操作 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]

Python/Numpy/Boolean Indexing: For each True value in array, modify the next 2 elements