python - 如果 cumsum 大于值，则重新启动 cumsum 并获取索引

Question

假设我有一个距离数组 x=[1,2,1,3,3,2,1,5,1,1]。

我想从 x 中获取索引，其中 cumsum 达到 10，在本例中，idx=[4,9]。

所以 cumsum 在满足条件后重新启动。

我可以用循环来完成，但是对于大型数组来说循环很慢，我想知道我是否可以用向量化的方式来完成。

Answer 1

一个有趣的方法

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]),)