使用NumPy实现数组元素上下左右数值求和的向量化方法

先初始化目标数组：

import numpy as np

# 创建示例数组
arr = np.array([
    [100., 100., 100., 100., 100.],
    [100.,   0.,   0.,   0., 100.],
    [100.,   0.,   0.,   0., 100.],
    [100.,   0.,   0.,   0., 100.],
    [100., 100., 100., 100., 100.]
])

利用NumPy的切片特性直接实现向量化计算，无需循环：

# 定位中间3×3区域
middle = arr[1:4, 1:4]

# 分别取上下左右四个方向的对应区域切片
top = arr[0:3, 1:4]    # 中间区域每个元素的上方值
bottom = arr[2:5, 1:4] # 中间区域每个元素的下方值
left = arr[1:4, 0:3]   # 中间区域每个元素的左方值
right = arr[1:4, 2:5]  # 中间区域每个元素的右方值

# 替换中间区域为四个方向值的和
middle[:] = top + bottom + left + right

执行后数组变为：

[[100. 100. 100. 100. 100.]
 [100. 200. 200. 200. 100.]
 [100. 200. 200. 200. 100.]
 [100. 200. 200. 200. 100.]
 [100. 100. 100. 100. 100.]]

原理说明 #

• 切片arr[1:4,1:4]精准定位中间3×3的零值区域，NumPy切片是原数组的视图，直接赋值会修改原数组
• 上下左右的切片维度都和中间区域一致（3×3），对应位置元素一一匹配，直接相加就是每个元素四个方向的和，完全符合向量化操作的要求

内容的提问来源于stack exchange，提问作者jianmin-chen