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

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 循环之类的来做到这一点，但必须有一种更 Python 的方式和更有效的方式！

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

这是做什么的：

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）
2. np.ufunc.reduceat(a, c) reduces a 由ufunc（在本例中为add）超出c[i]:c[i+1] 指定的范围。当i+1 溢出c 时，它会改为reduces 超过c[i]:-1
3. reduce 只是将数组压缩为单个值。例如，np.add.reduce(a) 等价于（但慢于）np.sum(a)（又慢于a.sum()）。但是，由于reduceat 将@jdehsa 的答案中的for 循环从python 推到numpy 核心编译的c 代码中，因此速度要快得多。

速度测试：

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

还有一个额外的好处是不浪费内存创建a的临时split副本

Answer 3

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

>>> 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的解决方案快一倍左右。