1

我有一个 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

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

所需的掩码输出：

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

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

4

1 回答 1

3

您可以使用它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]
于 2021-04-13T16:37:31.313 回答