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?
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]])
n=4 but this solution gives an error in that case. Can that be fixed?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.
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]]]
n instead of 2 :-) I was faster this time :-P