优化为Numpy数组非零值生成邻域缓冲区的函数性能

需求说明 #

实现函数：若原数组某位置的8邻域（水平、垂直、对角线方向）存在非零值，则返回数组对应位置为1，否则为0。

示例输入：

array([[10,  0,  0,  0,  0],
       [ 0,  0,  2,  0,  0],
       [ 0,  0,  0,  0,  0]])

示例输出：

array([[1, 1, 1, 1, 0],
       [1, 1, 1, 1, 0],
       [0, 1, 1, 1, 0]])

原实现的性能问题 #

原方案通过手动偏移数组并取最大值实现，但大量使用Python列表推导和循环，在处理大型数组时效率极低——Python层面的循环操作无法利用Numpy的矢量化优化优势，内存开销和计算时间都会显著增加。

高效优化方案 #

方案1：基于Scipy形态学膨胀（推荐） #

形态学膨胀操作天然匹配需求：将二值化后的数组用3x3全1结构元素膨胀，即可将所有非零值的8邻域标记为1。该操作底层为C实现，完全矢量化，性能拉满。

import numpy as np
from scipy.ndimage import binary_dilation

def expand_nonzero_neighbors(arr):
    # 数组二值化：非零转1，零保持0
    binary_arr = np.where(arr > 0, 1, 0)
    # 定义3x3结构元素，覆盖8邻域
    struct_elem = np.ones((3, 3), dtype=int)
    # 执行膨胀操作并转为int类型
    return binary_dilation(binary_arr, structure=struct_elem).astype(int)

# 测试示例
arr1 = np.array([[10, 0, 0, 0, 0],
                 [0, 0, 2, 0, 0],
                 [0, 0, 0, 0, 0]])
res = expand_nonzero_neighbors(arr1)
print(res)

方案2：纯Numpy实现（无Scipy依赖） #

利用Numpy的stride_tricks创建滑动窗口，计算每个3x3窗口的最大值，同样实现矢量化计算：

import numpy as np

def expand_nonzero_neighbors_np(arr):
    # 数组二值化
    binary_arr = np.where(arr > 0, 1, 0)
    # 给数组加一圈0填充，避免边界处理问题
    padded = np.pad(binary_arr, pad_width=1, mode='constant', constant_values=0)
    # 通过stride创建3x3滑动窗口视图（无额外内存开销）
    stride = padded.strides
    windows = np.lib.stride_tricks.as_strided(
        padded,
        shape=(binary_arr.shape[0], binary_arr.shape[1], 3, 3),
        strides=stride + stride
    )
    # 计算每个窗口的最大值，得到最终结果
    return windows.max(axis=(2, 3)).astype(int)

# 测试示例
arr1 = np.array([[10, 0, 0, 0, 0],
                 [0, 0, 2, 0, 0],
                 [0, 0, 0, 0, 0]])
res = expand_nonzero_neighbors_np(arr1)
print(res)

性能对比 #

• 原实现：处理1000x1000数组需数百毫秒甚至秒级（Python循环瓶颈）
• 优化方案：处理同规模数组仅需数毫秒（矢量化C底层计算）

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