在数组中找到所有元素之和小于某个限制的索引，快速进行。

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个值。那么瓶颈在哪里？如何加速此代码？

- lyvic

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个元素且limit = 100000.0/2的large数据集：

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

- YXD

---

对于大数组的情况，如果限制条件早早地达到了，这种方法会比较慢，因为它需要在测试限制条件之前计算整个累积和，对吧？ - M4rtini

是的，但即使有100000个元素，它仍然非常快-请参见更新。我个人不会费心去进一步优化，除非这是真正的瓶颈。 - YXD

使用Numba和装饰Python循环在这种情况下非常简单，并且效果很好。在100000个元素的情况下，比NumPy快近30倍。 - M4rtini

我喜欢这种方式的简洁性。性能足够好，技术非常清晰。cumsum() 是最大的区别。 - lyvic

@M4rtini 非常快！ - YXD

Answer 2

使用Numba，您可以显著加速Python循环。

import numba
import numpy as np

def numpyloop(large,limit):
    return np.argmin(large.cumsum() < limit)

@numba.autojit
def pythonloop(large,limit):
    sm = 0
    idx = 0
    for i in range(len(large)):
    #for i in range(large.shape[0]):
        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
    return idx

large = np.array([3] * 8000)
limit = 23996  

%timeit pythonloop(large,limit)
%timeit numpyloop(large,limit)

large = np.array([3] * 100000)
limit = 100000/2

%timeit pythonloop(large,limit)
%timeit numpyloop(large,limit)

Python：100000 次循环，3 次中的最佳表现：6.63 微秒每次循环
Numpy：10000 次循环，3 次中的最佳表现：33.2 微秒每次循环

大数组，小限制
Python：100000 次循环，3 次中的最佳表现：12.1 微秒每次循环
Numpy：1000 次循环，3 次中的最佳表现：351 微秒每次循环

- M4rtini