import Numpy as np

a = np.array([[5, 1, 8, 1, 6, 1, 3, 2],[2, 3, 4, 1, 6, 1, 4, 2]])
n = 2
[(a[0:2, i:i n]).sum(axis=1) for i in range(0,a.shape[1],n)]

The output is:

[array([6, 5]), array([9, 5]), array([7, 7]), array([5, 6])]

How can I get a 2D array instead of 4 arrays I got in above output...Is there a better more elegant way doing this using reshape?

CodePudding user response：

You can use sliding_window_view:

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

a = np.array([[5, 1, 8, 1, 6, 1, 3, 2], [2, 3, 4, 1, 6, 1, 4, 2]])
n = 2

out = sliding_window_view(a, (n, n)).sum(axis=n 1)[:, ::n].squeeze()

Output:

array([[6, 5],
       [9, 5],
       [7, 7],
       [5, 6]])

CodePudding user response：

You can always stack the results to get a single array:

import numpy as np

a = np.array([[5, 1, 8, 1, 6, 1, 3, 2],[2, 3, 4, 1, 6, 1, 4, 2]])
n = 2
np.stack([(a[0:2, i:i n]).sum(axis=1) for i in range(0,a.shape[1],n)])

Giving you:

array([[6, 5],
       [9, 5],
       [7, 7],
       [5, 6]])

Alternatively, you can reshape the array and sum:

a.T.reshape(-1, n, 2).sum(axis=1)

Giving the same result. But note, the array will need to be reshape-able to the given n, so n=4 is fine n=3 is an error.

CodePudding user response：

Another solution:

from numpy.lib.stride_tricks import sliding_window_view

a = np.array([[5, 1, 8, 1, 6, 1, 3, 2],[2, 3, 4, 1, 6, 1, 4, 2]])
n = 2

print(sliding_window_view(a, (n,n))[:,::n].sum(axis=n 1))

Prints:

[[[6 5]
  [9 5]
  [7 7]
  [5 6]]]