Python 按元素编号范围定义的步骤对 2D 数组求和

Question

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

输出为:

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

如何获得 2D 数组而不是上面输出中得到的 4 个数组...是否有更好更优雅的方法使用 reshape 来实现此目的？

Answer 1

您可以使用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()

输出:

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