python - Numpy 在索引周围创建掩码数组的最快方法

Question

我有一个一维数组:

arr = np.array([0, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 0, 3, 0, 0, 2, 0, 0, 0, 0, 0, 2, 3, 0, 0, 1, ...], dtype='uint16')

我想创建一个 mask 数组，该数组在任何大于 2 的值的 +/- N 索引处都为 True，产生以下结果(对于 N=3)

mask = [F, T, T, T, T, T, T, T, F, F, T, T, T, T, T, T, T, F, F, F, T, T, T, T, T, T, T, ...]

(注意，为了便于阅读，我使用了 T/F)

我需要这个速度相当快，因为​​我的实际数组有数百万个点长，N 可能约为 500。

编辑:similar problem

Answer 1

找到大于2的元素，然后将它们周围的元素设置为True:

a = np.array([0, 0, 0, 0, 3, 0, 0, 0, 1, 0, 0, 0, 0, 3, 0, 0, 2, 0])

N = 3
mask = a > 2
center = np.where(mask)[0]
mask[np.maximum(np.ravel(center - np.arange(1, 1 + N).reshape(-1, 1)), 0)] = True
mask[np.minimum(np.ravel(center + np.arange(1, 1 + N).reshape(-1, 1)), len(a)-1)] = True

感谢@Michael Szczesny 指出了边缘情况。 最大值和最小值确保索引不会(无意中)超出范围。