I need to do something like this
import numpy as np a = np.random.rand(1000) a -= .55 a[0] = 1 b = 0 for i in range(len(a)): b += a[i] if b < 0: print(i) break a lot, and preferably it should be swift. Can NumPy help me with that? I understand that NumPy is built for vector calculation and not for this. Still, it can calculate the sum of an array blazingly fast - can I give it this specific condition (the sum is negative for the first time) to stop and tell me the index number?
You can use numpy.cumsum() and numpy.argmax(). First, compute the cumulative sum of the elements. Then return True/False that elements < 0 and return the first index that is True for element < 0 with argmax().
>>> (a.cumsum() < 0).argmax() check both codes:
import numpy as np a = np.random.rand(1000) a -= .55 a[0] = 1 def check_1(a): b = 0 for i in range(len(a)): b += a[i] if b < 0: print(i) break def check_2(a): return (a.cumsum() < 0).argmax() Output: (Generate base random input)
>>> check_1(a) 6 >>> check_2(a) 6 
%timeit check_1(a) or %timeit check_2(a) in jupyter. for np.random.rand(1000) I got 6.81 µs ± 628 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) for both but for larger array your code is running faster.