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

Question

Lets say that n is a variable and I use a simple matrix example.

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]])

So desired_output[i] = sum of elements in the interval matrix[i - n, i + n] both inclusive. Is there a way to do this using NumPy and without python iteration? The numbers in the array can be arbitrary.

Answer 1

You can use one of my absolute favorite things to be added to numpy in 1.20: 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

First, apply the sliding window view:

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

Then sum it on the last axis:

arr = arr.sum(axis=2)

This gives you:

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

To the best of my knowledge there's no integer NaN in numpy so your integer/nan output is impossible. If you want floats you can pad easily with nans though:

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]])