通过滑动窗口重复 Numpy 数组

Question

来自以下形状数组(6, 3)：

>>> arr
[
    [1, 0, 1],
    [0, 0, 2],
    [1, 2, 0],
    [0, 1, 3],
    [2, 2, 1],
    [2, 0, 2]
]

我想根据n=4的滑动窗口重复这些值，给出一个形状为(6-n-1, n, 3)的新数组：

>>> new_arr
[
    [
        [1, 0, 1],
        [0, 0, 2],
        [1, 2, 0],
        [0, 1, 3]
    ],
    [
        [0, 0, 2],
        [1, 2, 0],
        [0, 1, 3],
        [2, 2, 1]
    ],
    [
        [1, 2, 0],
        [0, 1, 3],
        [2, 2, 1],
        [2, 0, 2]
    ]
]

使用循环相对简单，但在初始数组中有几百万个值（而不是本例中的 6 个）时会变得非常慢。

有没有更快的方法来使用 Numpy 原语到达new_arr？

Answer 1

您可以使用 NumPy，特别是 this 函数（仅 NumPy >= 1.20.0）：

from numpy.lib.stride_tricks import sliding_window_view

new_arr = sliding_window_view(arr, (n, arr.shape[1])).squeeze()