python - 选择索引周围的 NumPy 值

Question

我有两个 NumPy 数组:

import numpy as np

m = 3
x = np.array([1, 0, 0, np.inf, 0, 0, 1, 1, 2, np.inf, np.inf, 2, 3, 4, 5, 6, 7, 8, 9, 10])
y = np.arange(x.shape[0]-m+1)

假设 x 中只要存在 np.inf，该索引位置就称为 i。对于每个i，我想设置y[i-m+1:i+m] = np.inf的值。因此，替换后，y 应该如下所示:

array([0, np.inf, np.inf, np.inf, np.inf, np.inf, 6, np.inf, np.inf, np.inf, np.inf, np.inf, np.inf, 13, 14, 15, 16, 17])

当 m 的值增加或减少时，这也应该起作用。

Answer 1

我能想到的最简单的解决方案是使用 np.convolve函数来扩大掩模。这可以按如下方式完成:

mask = np.convolve(x==np.inf, [True]*(m*2-1), mode='same')
y[mask[:-m+1]] = np.inf