带复位条件的累计和

Question

嗨，所以我有一个数据框，我想在列的累积总和之一等于阈值时找到索引。然后它将重置并再次启动 cumsum。

例如：

    d = np.random.randn(10, 1) * 2
    df = pd.DataFrame(d.astype(int), columns=['data'])
    pd.concat([df,df.cumsum()],axis=1)

输出：

输出[34]：

   data  data1
0     1     1
1     2     3
2     3     6
3     2     8
4     0     8
5     1     9
6     0     9
7    -1     8
8     1     9
9     2    11

所以在上面的示例数据中，data1是第1列的cumsum。如果我设置thres=5这意味着每当第1列的运行总和大于或等于5时，我都会保存索引。之后，运行总和将重置并重新开始，直到达到下一个运行总和大于或等于 5。

现在我正在做一个循环并保持跟踪运行总和并手动重置。我想知道 pandas 中是否有一种快速的矢量化方法可以做到这一点，因为我的数据框有数百万行长。

谢谢

Answer 1

我对 pandas 不熟悉，但我的理解是它基于 numpy。使用 numpy，您可以定义可与累积一起使用的自定义函数。

这是我认为与您正在寻找的内容接近的一个：

import numpy as np
def capsum(array,cap):
   capAdd = np.frompyfunc(lambda a,b:a+b if a < cap else b,2,1)
   return capAdd.accumulate(values, dtype=np.object)

values = np.random.rand(1000000) * 3 // 1

result = capsum(values,5)  # --> produces the result in 0.17 sec.

我相信（或希望）您可以在数据帧上使用 numpy 函数。