Restart cumsum and get index if cumsum more than value

Question

Say I have an array of distances x=[1,2,1,3,3,2,1,5,1,1].

I want to get the indices from x where cumsum reaches 10, in this case, idx=[4,9].

So the cumsum restarts after the condition are met.

I can do it with a loop, but loops are slow for large arrays and I was wondering if I could do it in a vectorized way.

Answer 1

A fun method

sumlm = np.frompyfunc(lambda a,b:a+b if a < 10 else b,2,1)
newx=sumlm.accumulate(x, dtype=np.object)
newx
array([1, 3, 4, 7, 10, 2, 3, 8, 9, 10], dtype=object)
np.nonzero(newx==10)

(array([4, 9]),)