NumPy：查找运行总计超过值的索引

Question

我有一系列数字：

>>> x
array([20, 20,  1,  3, 13, 20, 25, 20, 10,  9, 20])

我想找到x的运行总数超过15的索引，每次超过15时总重置。因此，对于上面的示例数据，它应该给出：

>>> tot_indices(x)
array([0, 1, 4, 5, 6, 7, 9, 10])

目前，我正在这样做：

cumulative_total = np.copy(x)
ii = 0
cum_tot = 0
while ii < len(accum_spaces):
  cum_tot += x[ii]
  cumulative_total[ii] = cum_tot
  if cum_tot >= 15:
    cum_tot = 0
  ii += 1
indices = x[np.where(cum_tot >= 15)]

这样可行，但是NumPy是否有内置的方法来避免Python循环？

Answer 1

您可以使用numpy.cumsum

简化代码

cumulative_total = np.cumsum(x)
indices = x[np.where(cumulative_total >= 15)]