有没有办法仅使用 NumPy 来查找 2D 数组中最后 n 个元素和前 n 个元素的总和？

Question

假设n是一个变量，我使用一个简单的矩阵示例。

import numpy as np

n = 2

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

desired_output = np.array([[nan, nan, 10, 15, 20, 25, 30, 35, nan, nan], 
                           [nan, nan, 10, 15, 20, 25, 30, 35, nan, nan]])

所以desired_output[i] = 区间矩阵[i - n, i + n]中的元素之和，包括两者。 有没有办法使用 NumPy 而没有 python 迭代来做到这一点？ 数组中的数字可以是任意的。

Answer 1

您可以使用我最喜欢的东西之一在 1.20 中添加到 numpy： numpy.lib.stride_tricks.sliding_window_view

import numpy as np

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

首先，应用滑动 window 视图：

window_size = 2 * n + 1
arr = np.lib.stride_tricks.sliding_window_view(matrix, window_size, axis=1)

然后在最后一个轴上求和：

arr = arr.sum(axis=2)

这给了你：

>>> arr
array([[10, 15, 20, 25, 30, 35],
       [10, 15, 20, 25, 30, 35]])

据我所知，numpy 中没有 integer NaN 所以你的整数/nan output 是不可能的。 如果你想要浮动，你可以很容易地用 nans 填充：

arr_padded = np.full((arr.shape[0], arr.shape[1] + 2 * n), np.nan)
arr_padded[:, n:-1 * n] = arr

>>> arr_padded
array([[nan, nan, 10., 15., 20., 25., 30., 35., nan, nan],
       [nan, nan, 10., 15., 20., 25., 30., 35., nan, nan]])