子阵列的总和

时间:2014-01-19 18:38:48

标签: python optimization numpy

我有一个2d整数数组,我想总结它的2d子数组。两个数组都可以有任意维度,但我们可以假设子数组的数量级小于总数组。

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数组cumsumr_可以工作,我可以用一点循环来实现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.]]

3 个答案:

答案 0 :(得分:4)

有一个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”数组,就可以通过重新整形来获得所需的结果(因此一个块中的数字沿着单个轴排列),然后在该轴上调用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而不是重塑:

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

答案 1 :(得分:1)

代码中的更改是使用切片并使用sum()方法执行子数组的总和:

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(len(arr) // l):
        for j in range(len(arr[0]) // m):
            result[i, j] = arr[i*m:(i+1)*m, j*l:(j+1)*l].sum()
    return result

执行一些非常简单的基准测试表明2x2案例中的更慢,大约相当于3x3案例中的方法,更大的子阵列更快(sub_sums2您的版本的代码):

In [19]: arr = np.asarray([range(100)] * 100)

In [20]: %timeit sub_sums(arr, 2, 2)
10 loops, best of 3: 21.8 ms per loop

In [21]: %timeit sub_sums2(arr, 2, 2)
100 loops, best of 3: 9.56 ms per loop

In [22]: %timeit sub_sums(arr, 3, 3)
100 loops, best of 3: 9.58 ms per loop

In [23]: %timeit sub_sums2(arr, 3, 3)
100 loops, best of 3: 9.36 ms per loop

In [24]: %timeit sub_sums(arr, 4, 4)
100 loops, best of 3: 5.58 ms per loop

In [25]: %timeit sub_sums2(arr, 4, 4)
100 loops, best of 3: 9.56 ms per loop

In [26]: %timeit sub_sums(arr, 10, 10)
1000 loops, best of 3: 939 us per loop

In [27]: %timeit sub_sums2(arr, 10, 10)
100 loops, best of 3: 9.48 ms per loop

请注意,对于10x10个子数组,它的 1000 时间更快。在2x2情况下,它的速度大约是其两倍。你的方法基本上总是在同一时间,而我的实现随着更大的子阵列变得更快。

我很确定我们可以避免显式使用for循环(可能重新整形数组,以便它将子数组作为行?),但我不是{{1}的专家我可能需要一些时间才能找到最终解决方案。但是我认为3个数量级已经是一个很好的改进。

答案 2 :(得分:1)

这是更简单的方法:

In [160]: import numpy as np

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

In [162]: np.add.reduceat(arr, [0], axis=1)
Out[162]: 
array([[10],
       [30],
       [50],
       [70]])

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

In [164]: import numpy as np

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

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

In [167]: np.add.reduceat(arr, [0], axis=1)
Out[167]: 
array([[10],
       [30],
       [50],
       [70]])