加快基于组中另一列的选择性累加和

Question

我有一个数据框，我想要 groupby 2 列，然后创建一个新列，该列将具有第 3 列的累计总和，其中计数取决于第四列的值。我有可用的代码，但速度非常慢。我如何加快速度？

所以在下面的例子中，如果 dir 等于 up 乘以 date 和 sym，我想要 Qty 的累加和。

在 R 和 data.table 中，这将是一个简单的单行代码，可以非常快速地完成执行：

d1[,newColName:=cumsum(Qty*(dir=="up")),by=c("date","sym")]

我在 python 和 pandas 中提出的是一个非常慢（但有效）的函数，具有以下用途：

def test(x):
  return pd.Series([ a*b for a,b in zip([ 1 if y == "up" else 0 for y in x["dir"] ], x["Qty"].tolist()) ]).cumsum()

# example use
d1[1:20].groupby(["date","sym"])[["dir","Qty"]].apply(test) # too slow to run over he whole data set

数据块示例：

d1[["date","sym","dir","Qty" ]]
Out[102]: 
             date sym   dir  Qty
0      2019-10-29  A1    up    9
1      2019-10-29  A1  down    1
2      2019-10-29  A1  down   11
3      2019-10-29  A1    up    2
4      2019-10-29  A1    up    3

我如何才能加快速度，以便我在 python 中实际运行处理大量数据？它不一定是 pandas 顺便说一句，但应该是 python.

所以这是我想要得到的输出：

> d1
             date sym  dir Qty newColName
    1: 2019-10-29  A1   up   9          9
    2: 2019-10-29  A1 down   1          9
    3: 2019-10-29  A1 down  11          9
    4: 2019-10-29  A1   up   2         11
    5: 2019-10-29  A1   up   3         14

Answer 1

您可以尝试以下方法：

>>> df.loc[df.dir.eq('up'), 'newColName'] = df[df.dir.eq('up')].groupby(['date', 'sym'])['Qty'].cumsum()
>>> df['newColName'] = df['newColName'].ffill(downcast='infer')
>>> df
         date sym   dir  Qty  newColName
0  2019-10-29  A1    up    9           9
1  2019-10-29  A1  down    1           9
2  2019-10-29  A1  down   11           9
3  2019-10-29  A1    up    2          11
4  2019-10-29  A1    up    3          14

或者，使用与 reindex 相同的东西：

>>> df['newColName'] = (df[df.dir.eq('up')]
         .groupby(['date', 'sym'])['Qty']
         .cumsum().reindex(df.index)
         .ffill(downcast='infer')
    )
>>> df
         date sym   dir  Qty  newColName
0  2019-10-29  A1    up    9           9
1  2019-10-29  A1  down    1           9
2  2019-10-29  A1  down   11           9
3  2019-10-29  A1    up    2          11
4  2019-10-29  A1    up    3          14

加快基于组中另一列的选择性累加和

Speed up selective cumulative sum based on another column in a group

python

r

dataframe

pandas

cumsum