首页 > 解决方案 > Pandas,如何通过 groupby() / 条件函数调用应用更改数据框?

问题描述

我有一个超过 1m 行的巨大时间序列数据框。它有每月股票收益列表,我想创建一个新行来跟踪前 3 个月的滚动总和。数据框首先包含所有公司 A 行,然后是所有公司 B 行,然后是所有公司 C 行......

例如:

          date                COMNAM    PRC    RET
395 2017-02-28       GAS NATURAL INC 12.650  0.000
396 2017-03-31       GAS NATURAL INC 12.700  0.010
397 2017-04-28       GAS NATURAL INC 12.500 -0.016
398 2017-05-31       GAS NATURAL INC 12.700  0.016
399 2017-06-30       GAS NATURAL INC 12.925  0.024
400 2017-07-31       GAS NATURAL INC 12.950  0.002
401 2017-08-31       GAS NATURAL INC    nan    nan
402 1985-12-31                   NaN    nan    nan
403 1986-01-31  MOBILE NATIONAL CORP 11.625    nan
404 1986-02-28  MOBILE NATIONAL CORP 13.250  0.140
405 1986-03-31  MOBILE NATIONAL CORP 14.188  0.071
406 1986-04-30  MOBILE NATIONAL CORP 14.938  0.053
407 1986-05-30  MOBILE NATIONAL CORP 14.625 -0.021
408 1986-06-30  MOBILE NATIONAL CORP 12.688 -0.132
409 1986-07-31  MOBILE NATIONAL CORP 13.312  0.049
410 1986-08-29  MOBILE NATIONAL CORP 13.312  0.000
411 1986-09-30  MOBILE NATIONAL CORP 14.250  0.070
412 1986-10-31  MOBILE NATIONAL CORP 13.375 -0.061
413 1986-11-28  MOBILE NATIONAL CORP 13.375  0.000
414 1986-12-31  MOBILE NATIONAL CORP 12.375 -0.075

rolling() 函数将为我提供前 3 个月的总和,但这将包括每家公司第一个日期的前一个股票的最后回报。我觉得 groupby() 函数可能会有所帮助,但我有点坚持如何做。还是我想多了,有更好的方法我什至不需要 groupby?

标签: pandastime-seriespandas-groupby

解决方案


要计算前 3 个月的滚动总和包括当前月份),对于感兴趣的两列,从当前行组中定义以下函数:

def mySum(grp):
    return grp[['PRC', 'RET']].shift().rolling(3).sum() 

然后,要获得每个组(公司)的此类滚动总和,请运行:

result = df.join(df[df.COMNAM.notnull()].groupby('COMNAM').apply(mySum)\
    .reset_index(level=0, drop=True).add_prefix('r'))

结果是当前df与为每个组(公司)调用上述函数的结果之间的连接。中间结果的列名前面带有r,以标记滚动和。

对于您的数据样本,结果是:

         date                COMNAM     PRC    RET    rPRC   rRET
0  2017-02-28       GAS NATURAL INC  12.650  0.000     NaN    NaN
1  2017-03-31       GAS NATURAL INC  12.700  0.010     NaN    NaN
2  2017-04-28       GAS NATURAL INC  12.500 -0.016     NaN    NaN
3  2017-05-31       GAS NATURAL INC  12.700  0.016  37.850 -0.006
4  2017-06-30       GAS NATURAL INC  12.925  0.024  37.900  0.010
5  2017-07-31       GAS NATURAL INC  12.950  0.002  38.125  0.024
6  2017-08-31       GAS NATURAL INC     NaN    NaN  38.575  0.042
7  1985-12-31                   NaN     NaN    NaN     NaN    NaN
8  1986-01-31  MOBILE NATIONAL CORP  11.625    NaN     NaN    NaN
9  1986-02-28  MOBILE NATIONAL CORP  13.250  0.140     NaN    NaN
10 1986-03-31  MOBILE NATIONAL CORP  14.188  0.071     NaN    NaN
11 1986-04-30  MOBILE NATIONAL CORP  14.938  0.053  39.063    NaN
12 1986-05-30  MOBILE NATIONAL CORP  14.625 -0.021  42.376  0.264
13 1986-06-30  MOBILE NATIONAL CORP  12.688 -0.132  43.751  0.103
14 1986-07-31  MOBILE NATIONAL CORP  13.312  0.049  42.251 -0.100
15 1986-08-29  MOBILE NATIONAL CORP  13.312  0.000  40.625 -0.104
16 1986-09-30  MOBILE NATIONAL CORP  14.250  0.070  39.312 -0.083
17 1986-10-31  MOBILE NATIONAL CORP  13.375 -0.061  40.874  0.119
18 1986-11-28  MOBILE NATIONAL CORP  13.375  0.000  40.937  0.009
19 1986-12-31  MOBILE NATIONAL CORP  12.375 -0.075  41.000  0.009

如果要“忽略” NaN值(将它们视为0),请将函数更改为:

def mySum(grp):
    return grp[['PRC', 'RET']].fillna(0).shift().rolling(3).sum()

推荐阅读