python - 子数组的总和

Question

我有一个二维整数数组，我想求和它的二维子数组。两个数组都可以具有任意维度，尽管我们可以假设子数组将比总数组小几个数量级。

python 中的引用实现很简单:

def sub_sums(arr, l, m):
    result = np.zeros((len(arr) // l, len(arr[0]) // m))
    rows = len(arr) // l * l
    cols = len(arr[0]) // m * m
    for i in range(rows):
        for j in range(cols):
            result[i // l, j // m] += arr[i, j]
    return result

问题是我如何使用 numpy 做到最好，希望在 python 中完全没有任何循环。对于 1d 数组 cumsum 和 r_ 可以工作，我可以使用它和一些循环来实现 2d 的解决方案，但我仍在学习 numpy 并且我正在几乎可以肯定有一些更聪明的方法。

示例输出:

arr = np.asarray([range(0, 5),
                  range(4, 9),
                  range(8, 13),
                  range(12, 17)])
result = sub_sums(arr, 2, 2)

给出:

[[ 0  1  2  3  4]
 [ 4  5  6  7  8]
 [ 8  9 10 11 12]
 [12 13 14 15 16]]

[[ 10.  18.]
 [ 42.  50.]]

Answer 1

有一个blockshaped function它做的事情非常接近你想要的:

In [81]: arr
Out[81]: 
array([[ 0,  1,  2,  3,  4],
       [ 4,  5,  6,  7,  8],
       [ 8,  9, 10, 11, 12],
       [12, 13, 14, 15, 16]])

In [82]: blockshaped(arr[:,:4], 2,2)
Out[82]: 
array([[[ 0,  1],
        [ 4,  5]],

       [[ 2,  3],
        [ 6,  7]],

       [[ 8,  9],
        [12, 13]],

       [[10, 11],
        [14, 15]]])

In [83]: blockshaped(arr[:,:4], 2,2).shape
Out[83]: (4, 2, 2)

一旦你有了“blockshaped”数组，你就可以通过整形(所以一个 block 中的数字沿着一个轴排列)然后调用 sum 方法来获得所需的结果轴。

因此，通过对 blockshaped 函数稍作修改，您可以像这样定义 sub_sums:

import numpy as np

def sub_sums(arr, nrows, ncols):
    h, w = arr.shape
    h = (h // nrows)*nrows
    w = (w // ncols)*ncols
    arr = arr[:h,:w]
    return (arr.reshape(h // nrows, nrows, -1, ncols)
               .swapaxes(1, 2)
               .reshape(h // nrows, w // ncols, -1).sum(axis=-1))

arr = np.asarray([range(0, 5),
                  range(4, 9),
                  range(8, 13),
                  range(12, 17)])

print(sub_sums(arr, 2, 2))

产量

[[10 18]
 [42 50]]

---

编辑:Ophion 提供了一个很好的改进——在求和之前使用 np.einsum 而不是 reshape :

def sub_sums_ophion(arr, nrows, ncols):
    h, w = arr.shape
    h = (h // nrows)*nrows
    w = (w // ncols)*ncols
    arr = arr[:h,:w]
    return np.einsum('ijkl->ik', arr.reshape(h // nrows, nrows, -1, ncols))

---

In [105]: %timeit sub_sums(arr, 2, 2)
10000 loops, best of 3: 112 µs per loop

In [106]: %timeit sub_sums_ophion(arr, 2, 2)
10000 loops, best of 3: 76.2 µs per loop