如何快速为N×N二维numpy数组生成带边界环绕的邻域张量？

方法一：广播索引直接提取（直观易维护） #

无需额外构建填充数组，直接通过广播生成所有邻域的索引，结合取模运算实现边缘环绕，全程矢量化操作，性能远高于逐一遍历。

示例代码：

import numpy as np

def get_circular_neighbors(A, size):
    N = A.shape[0]
    # 生成所有中心位置的网格坐标
    i, j = np.meshgrid(np.arange(N), np.arange(N), indexing='ij')
    # 生成邻域的偏移范围（从 -size 到 size）
    di = np.arange(-size, size + 1)
    dj = np.arange(-size, size + 1)
    # 广播计算所有窗口的行/列索引，取模N实现循环环绕
    rows = (i[..., None, None] + di[None, None, :, None]) % N
    cols = (j[..., None, None] + dj[None, None, None, :]) % N
    # 一次性提取所有邻域窗口
    return A[rows, cols]

用法示例：

# 测试用例
N = 5
size = 1
A = np.arange(N*N).reshape(N, N)
neighbors = get_circular_neighbors(A, size)
print(neighbors.shape)  # 输出 (5,5,3,3)

方法二：as_strided滑动窗口视图（性能最优） #

如果追求极致性能，可使用numpy.lib.stride_tricks.as_strided直接创建内存视图，完全避免数据复制，速度最快。步骤为先构建循环填充的扩展数组，再通过步长操作生成所有窗口。

示例代码：

import numpy as np
from numpy.lib.stride_tricks import as_strided

def get_circular_neighbors_fast(A, size):
    N = A.shape[0]
    # 构建循环填充的扩展数组（边缘环绕）
    top = A[-size:, :]
    bottom = A[:size, :]
    left = A[:, -size:]
    right = A[:, :size]
    top_left = top[:, -size:]
    top_right = top[:, :size]
    bottom_left = bottom[:, -size:]
    bottom_right = bottom[:, :size]
    
    A_ext = np.vstack([
        np.hstack([top_left, top, top_right]),
        np.hstack([left, A, right]),
        np.hstack([bottom_left, bottom, bottom_right])
    ])
    # 计算步长：复用原数组的行/列步长，对应窗口维度
    strides = A_ext.strides + A_ext.strides
    # 目标形状：(N, N, 2*size+1, 2*size+1)
    shape = (N, N, 2*size + 1, 2*size + 1)
    # 创建滑动窗口视图
    return as_strided(A_ext, shape=shape, strides=strides)

注意：as_strided生成的是原数组的视图，修改返回的窗口会影响原数组。若需要独立副本，可在最后追加.copy()。

方法对比 #

• 方法一：代码简洁直观，内存占用小，适合绝大多数场景。
• 方法二：性能极致，适合超大规模数组（如N>1000），但需注意内存视图的副作用。

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