如何用Numpy高效生成每行固定步长偏移的批量移位数组？

推荐方法：使用numpy.lib.stride_tricks.sliding_window_view（Numpy 1.20+） #

这是Numpy官方提供的滑动窗口工具，无需手动处理内存，安全且高效，生成的是原数组的视图（非复制，节省内存），仅在修改时才会触发复制操作。

示例代码：

import numpy as np

# 原数组（示例：1到9的连续整数）
foo = np.arange(1, 10)
batch_size = 5

# 生成滑动窗口数组
result = np.lib.stride_tricks.sliding_window_view(foo, window_size=batch_size)

print(result)

输出结果：

[[1 2 3 4 5]
 [2 3 4 5 6]
 [3 4 5 6 7]
 [4 5 6 7 8]
 [5 6 7 8 9]]

旧版本Numpy兼容方案：numpy.lib.stride_tricks.as_strided #

如果你的Numpy版本低于1.20，可以用as_strided手动构造视图，但必须严格计算形状和步长，否则可能访问越界内存导致错误：

import numpy as np

foo = np.arange(1, 10)
batch_size = 5
num_rows = len(foo) - batch_size + 1

# 计算步长：原数组的元素步长，行和列的步长相同（因为偏移步长为1）
strides = (foo.strides[0], foo.strides[0])
result = np.lib.stride_tricks.as_strided(foo, shape=(num_rows, batch_size), strides=strides)

print(result)

特殊场景简化方案：广播（仅适用于连续规律数组） #

如果原数组是连续递增/递减的数值（比如示例中的1,2,3...），可以用广播直接生成结果，无需依赖滑动窗口工具：

import numpy as np

batch_size = 5
# 生成行索引和列索引，通过广播相加得到结果
rows = np.arange(batch_size)
cols = np.arange(batch_size)
result = rows[:, None] + cols + 1

print(result)

性能对比 #

• 你原来的循环实现是逐行赋值，时间复杂度为O(n*m)，且Python循环本身效率较低。
• 上述Numpy原生方法都是矢量化操作，由底层C实现，时间复杂度更低，尤其是处理大数组时性能提升非常明显。

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