Is there a way to find the sum of the last n and first n elements in a 2D array using only NumPy?-CodePudding

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.

CodePudding user response：

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