如何无需循环从NumPy数组生成切片数组？

嘿，很高兴帮你优化这段代码！用for循环实现虽然直观，但NumPy提供了更高效的矢量化方法来完成这个任务，下面给你介绍几种常用方案：

方法1：使用numpy.lib.stride_tricks.as_strided #

这是性能最优的方法之一，它直接利用数组的内存 stride 创建视图（不会复制原数据），运行效率极高：

import numpy as np

def calc_matrix(arr, k):
    n = len(arr) - k + 1
    # 定义新数组的形状和内存步长
    shape = (n, k)
    strides = (arr.strides[0], arr.strides[0])
    return np.lib.stride_tricks.as_strided(arr, shape=shape, strides=strides)

注意：该方法返回的是原数组的视图，修改返回数组会同步影响原数组。如果需要独立的数组副本，可以在最后加上.copy()。

方法2：利用广播生成索引矩阵 #

这种方法更直观，通过生成索引矩阵直接提取切片，不需要理解内存步长的概念：

import numpy as np

def calc_matrix(arr, k):
    n = len(arr) - k + 1
    # 生成每行的起始索引（转为列向量实现广播）
    row_starts = np.arange(n)[:, np.newaxis]
    # 生成切片内的偏移索引
    col_offsets = np.arange(k)
    # 广播相加得到所有元素的索引
    all_indices = row_starts + col_offsets
    return arr[all_indices]

该方法会创建新的数组，与原数组内存独立，适合需要修改结果又不想影响原数据的场景。

方法3：使用numpy.lib.stride_tricks.sliding_window_view（NumPy 1.20+） #

如果你的NumPy版本在1.20及以上，官方提供了更安全的滑动窗口工具，它封装了as_strided的底层逻辑，使用更简单且不易出错：

import numpy as np

def calc_matrix(arr, k):
    return np.lib.stride_tricks.sliding_window_view(arr, window_shape=k)

默认返回原数组的视图，同样可以通过.copy()获取独立副本。

效果验证 #

你可以用以下代码测试所有方法的一致性：

arr = np.array([1,2,3,4,5])
k = 3
print(calc_matrix(arr, k))

输出均为：

[[1 2 3]
 [2 3 4]
 [3 4 5]]

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