python - 快速查找数组中所有元素的总和小于限制的索引

Question

给定一个大数组。我正在寻找数组中所有元素加起来小于 limit 的数字的索引。我找到了两种方法:

import time as tm
import numpy as nm

# Data that we are working with
large = nm.array([3] * 8000)
limit = 23996

# Numpy version, hoping it would be faster
start = tm.time()  # Start timing
left1 = nm.tril([large] * len(large))  # Build triangular matrix
left2 = nm.sum(left1, 1)  # Sum up all rows of the matrix
idx = nm.where(left2 >= limit)[0][0]  # Check what row exceeds the limit
stop = tm.time()
print "Numpy result :", idx
print "Numpy took :", stop - start, " seconds"

# Python loop
sm = 0  # dynamic sum of elements
start = tm.time()
for i in range(len(large)):
    sm += large[i]  # sum up elements one by one
    if sm >= limit:  # check if the sum exceeds the limit
        idx = i
        break  # If limit is reached, stop looping.
    else:
        idx = i
stop = tm.time()
print "Loop result :", idx
print "Loop took :", stop - start, " seconds"

不幸的是，如果数组大得多，则 numpy 版本会耗尽内存。更大的意思是 100 000 个值。当然，这给出了一个大矩阵，但是 for 循环需要 2 分钟。也可以遍历这 100 000 个值。那么，瓶颈在哪里呢？我怎样才能加快这段代码的速度？

Answer 1

您可以通过以下方式获得:

np.argmin(large.cumsum() < limit)

或等效

(large.cumsum() < limit).argmin()

在 IPython 中:

In [6]: %timeit (large.cumsum() < limit).argmin()
10000 loops, best of 3: 33.8 µs per loop

对于包含 100000 个元素的 large，并且 limit = 100000.0/2

In [4]: %timeit (large.cumsum() < limit).argmin()
1000 loops, best of 3: 444 µs per loop

它没有任何真正的区别，但通常import numpy as np 而不是 import numpy as nm。

文档:

• http://docs.scipy.org/doc/numpy/reference/generated/numpy.cumsum.html
• http://docs.scipy.org/doc/numpy/reference/generated/numpy.argmin.html