如何在NumPy中实现依赖邻域位置的函数向量化

这类需要访问相邻元素的邻域计算属于典型的二维卷积场景，直接用卷积操作即可实现完全向量化，无需手动编写循环。

最简实现方案（基于scipy卷积） #

该方案的边界处理逻辑和你提供的原生循环完全一致，边界外默认填充0：

import numpy as np
from scipy.ndimage import convolve

mask = np.array([[0,0,0,0],
                 [1,0,0,0],
                 [0,0,0,1], 
                 [1,0,0,0]])
# 定义3×3全1卷积核，作用是对每个位置的邻域求和
kernel = np.ones((3,3), dtype=np.int32)
# 卷积计算所有位置的3×3邻域和
neighbor_sum = convolve(mask, kernel, mode='constant', cval=0)
# 同时满足两个条件的位置赋值为1
out = np.zeros_like(mask, dtype=np.float64)
out[(neighbor_sum > 1) & (mask != 1)] = 1

运行输出和原生循环结果完全相同：

[[0. 0. 0. 0.]
 [0. 0. 0. 0.]
 [0. 1. 0. 0.]
 [0. 0. 0. 0.]]

NumPy原生实现方案（无第三方依赖） #

如果不想引入scipy，也可以用NumPy自带的滑动窗口视图实现：

import numpy as np

mask = np.array([[0,0,0,0],
                 [1,0,0,0],
                 [0,0,0,1], 
                 [1,0,0,0]])
# 先给原数组四周填充一圈0，处理边界邻域缺失问题
padded_mask = np.pad(mask, pad_width=1, mode='constant', constant_values=0)
# 生成所有3×3滑动窗口的视图
windows = np.lib.stride_tricks.sliding_window_view(padded_mask, window_shape=(3,3))
# 对每个窗口求和得到邻域和数组
neighbor_sum = windows.sum(axis=(2,3))
# 条件判断生成输出
out = np.zeros_like(mask, dtype=np.float64)
out[(neighbor_sum > 1) & (mask != 1)] = 1

两种方案均为完全向量化实现，数组规模越大，相比原生Python循环的性能优势越明显。

内容的提问来源于stack exchange，提问作者Jkind9