Question

给出以下数据（在python 2.7中）：

import numpy as np
a = np.array([1,2,3,4,5,6,7,8,9,10,11,12,14])
b = np.array([8,2,3])

我希望获得a中前8个元素的总和，然后得到9和10元素的总和，最后得到最后3个元素（基本b中的信息）。所需的输出是：

[36, 19, 37]

我可以用for循环这样做，但必须有更多的pythonic方式和更有效的方法！

Answer 1

使用np.split：

很容易

result = [part.sum() for part in np.split(a, np.cumsum(b))[:-1]]
print(result)
>>> [36, 19, 37]

Answer 2

比np.split更快的方法是：

np.add.reduceat(a, np.r_[0, np.cumsum(b)[:-1]])

这是做什么的：

创建一个与b对应的升序索引数组，对应于您想要求和的范围 - 为简单起见，您可以指定c = np.r_[0, np.cumsum(b)[:-1]]，例如array([0, 8, 10]) - 哪个是0除了b（np.cumsum(b) -> array([8, 10, 13])的累积总和的最后一个元素之外的所有元素（np.ufunc.reduceat的域都不包含端点，所以我们必须摆脱那13）
np.ufunc.reduceat(a, c) reduce a ufunc add（在这种情况下，c[i]:c[i+1]）超出i+1指定的范围。如果c溢出reduce，则会c[i]:-1超过reduce
np.add.reduce(a)只是将数组压缩为单个值。例如，np.sum(a)相当于（但慢于）a.sum()（反过来慢于reduceat）。但是，由于for将@jdehsa中的numpy循环从python推送到b = np.random.randint(1,10,(10000,)) a = np.random.randint(1,10,(np.sum(b),)) %timeit np.add.reduceat(a, np.r_[0, np.cumsum(b)[:-1]]) 1000 loops, best of 3: 293 µs per loop %timeit [part.sum() for part in np.split(a, np.cumsum(b))[:-1]] 10 loops, best of 3: 44.6 ms per loop核心编译的c代码中，因此速度要快得多。

速度测试：

split

还有一项额外的好处，就是不浪费内存来创建a的临时testcsv2副本

Answer 3

您可以使用reduceat的np.add ufunc方法。您只需要在索引前添加一个零并丢弃最后一个索引（如果它覆盖整个数组）：

>>> import numpy as np
>>> a = np.array([1,2,3,4,5,6,7,8,9,10,11,12,14])
>>> b = np.array([8,2,3])
>>> np.add.reduceat(a, np.append([0], np.cumsum(b)[:-1]))
array([36, 19, 37], dtype=int32)

[:-1]丢弃最后一个索引，np.append([0],在索引前面加零。

请注意，这是DanielFs answer的稍微改编的变体。

如果你不喜欢append，你也可以自己创建一个包含索引的新数组：

>>> b_sum = np.zeros_like(b)
>>> np.cumsum(b[:-1], out=b_sum[1:])  # insert the cumsum in the b_sum array directly
>>> np.add.reduceat(a, b_sum)
array([36, 19, 37], dtype=int32)

Answer 4

如果必须使用b：

中的元素，则可以这样做

import numpy as np
a = np.array([1,2,3,4,5,6,7,8,9,10,11,12,14])
b = np.array([8,2,3])

c = np.array([np.sum(a[:b[0]]),np.sum(a[b[0]:b[0]+b[1]]),np.sum(a[-b[2]:])])

Answer 5

使用numba的解决方案

@Daniel F已经有了一个很好的pythonic答案。我不想展示另一种不那么pythonic但更快的解决方案。您可以在Python中使用循环，但如果您不想获得合理的速度，则必须使用编译器。 Numba很容易使用，所以我不想在这里举个例子。

import numba as nb
import numpy as np
import time
def main():
    b = np.random.randint(1,10,(10000,))
    a = np.random.randint(1,10,(np.sum(b),))

    nb_splitsum = nb.njit(nb.int32[:](nb.int32[:], nb.int32[:]),nogil=True)(splitsum)

    t1=time.time()
    for i in xrange(0,1000):
        c=nb_splitsum(a,b)

    print("Numba Solution")
    print(time.time()-t1)

    t1=time.time()
    for i in xrange(0,1000):
        c=np.add.reduceat(a, np.r_[0, np.cumsum(b)[:-1]])
    print("Numpy Solution")
    print(time.time()-t1)

def splitsum(a,b):
    sum=np.empty(b.shape[0],dtype=np.int32)
    ii=0
    for i in range(0,b.shape[0]):
        for j in range(0,b[i]):
            sum[i]+=a[ii]
            ii+=1
    return sum

if __name__ == "__main__":
    main()


#Output
Numba Solution
0.125
Numpy Solution
0.280999898911

我的计算机上的编译开销约为0.15秒。但是在编译函数时，上面显示的解决方案的速度大约是纯numpy解决方案的两倍。

计算数组特定区域总和的最快方法

5 个答案: