如何快速获取大尺寸mxn NumPy数组中距离l内的相邻元素？

针对大尺寸m×n数组的邻域元素操作（比如图像像素邻域处理），避免逐元素遍历，用numpy的向量化运算可以大幅提升效率，以下是针对你需求的实现方案：

核心思路 #

以你要求的曼哈顿距离l=1（四邻域：上下左右）为例，通过数组滚动移位生成各方向的邻域矩阵，过滤掉边界外的无效元素后，再批量计算源元素与邻域元素的乘积。

代码实现 #

1. 导入依赖并定义原数组 #

import numpy as np

x = np.array([[1,2,3],
              [4,5,6], 
              [7,8,9]])

2. 生成四邻域矩阵 #

通过np.roll生成上下左右四个方向的邻域，同时用掩码标记边界外的无效元素：

# 定义四个方向的偏移：上、下、左、右
shifts = [(-1, 0), (1, 0), (0, -1), (0, 1)]
neighbors = []

for dy, dx in shifts:
    # 滚动移位得到邻域数组
    rolled = np.roll(np.roll(x, dy, axis=0), dx, axis=1)
    # 标记边界外的元素为无效（用NaN）
    mask = np.ones_like(x, dtype=bool)
    if dy == -1:
        mask[0, :] = False
    elif dy == 1:
        mask[-1, :] = False
    if dx == -1:
        mask[:, 0] = False
    elif dx == 1:
        mask[:, -1] = False
    rolled[~mask] = np.nan
    neighbors.append(rolled)

# 合并所有邻域为一个三维数组：(4, m, n)
neighbors_stack = np.stack(neighbors)

3. 提取每个元素的有效邻域 #

# 遍历每个位置，收集非NaN的邻域元素
y = []
for i in range(x.shape[0]):
    row = []
    for j in range(x.shape[1]):
        valid_neighbors = neighbors_stack[:, i, j][~np.isnan(neighbors_stack[:, i, j])].tolist()
        row.append(valid_neighbors)
    y.append(row)

# 输出邻域列表，和你示例的y一致
print(y)

4. 计算源元素与邻域元素的乘积 #

# 批量计算乘积，保持邻域结构
product_y = []
for i in range(x.shape[0]):
    row = []
    for j in range(x.shape[1]):
        src_val = x[i, j]
        products = [src_val * n for n in y[i][j]]
        row.append(products)
    product_y.append(row)

print(product_y)

性能优化说明 #

对于超大数组，上述最后一步的嵌套循环可以进一步用numpy向量化替代：通过广播机制直接计算乘积，再过滤NaN值。示例如下：

# 向量化计算乘积，避免嵌套循环
x_expanded = x[np.newaxis, :, :]
product_matrix = x_expanded * neighbors_stack
# 转换为列表结构，过滤NaN
product_y_vectorized = [[product_matrix[:, i, j][~np.isnan(product_matrix[:, i, j])].tolist() 
                         for j in range(x.shape[1])] 
                        for i in range(x.shape[0])]

这种方式完全依赖numpy的C底层运算，处理大数组时效率远高于纯Python遍历。

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