A fast numpy way to find index in array where cumulative sum becomes greater?

Question

Basically, the logic of my problem is:

running_sum = my_array.cumsum()
greater_than_threshold = running_sum > threshold
index = greater_than_threshold.searchsorted(True)

That is: Find the first index for which the cumulative sum of entries in my_array is above a threshold.

Now the problem is: I know that my_array will be large, but that the condition will be met fairly early. Of course that means I could just do a simple while loop to manually figure out when the cumulative sum is larger than the threshold, but I am wondering if there's a numpythonic way, i.e., a way to test for some condition without having the entire array evaluated.

Answer 1

EDIT: This method is slower than using NumPy's searchsorted and cumsum, see user2357112's comments and timeit test.

cumsum will calculate cumulative sums for the entire array. Instead, just iterate over the array yourself:

running_sum = 0
for index, entry in enumerate(my_array.flat):
    running_sum += entry
    if running_sum > threshold:
        break
if running_sum < threshold:
    index = -1 #if the sum never reaches the threshold