Find index in array to which the sum of all elements is smaller than a limit, quickly

Question

Given is a large array. I am looking for the index up to which all elements in the array add up to a number smaller than limit. I found two ways to do so:

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"

Unfortunately, the numpy version runs out of memory if the array is much larger. By larger I mean 100 000 values. Of course, this gives a big matrix, but the for-loop takes 2min. to run through those 100 000 values just as well. So, where is the bottleneck? How can I speed this code up?

Answer 1

You can get this with:

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

or equivalently

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

In IPython:

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

for large with 100000 elements, and limit = 100000.0/2

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

It does not make any real difference, but it is conventional to import numpy as np rather than import numpy as nm.

Documentation:

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

Answer 2

Using numba you can speed up the python loop significantly.

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 loops, best of 3: 6.63 µs per loop
Numpy: 10000 loops, best of 3: 33.2 µs per loop

Large array, small limit
Python: 100000 loops, best of 3: 12.1 µs per loop
Numpy: 1000 loops, best of 3: 351 µs per loop