首页 > 解决方案 > Pandas 获取每组未来 6 个月数据的列值总和

问题描述

我有一个df看起来像这样的数据框:

                         id               date             num
0                         1            2005-01-01           7
1                         1            2006-01-05           2
2                         2            2005-01-01           1
3                         2            2005-04-01           1
4                         3            2009-01-01           2

对于每个id组,我想获得num未来 6 个月数据的总和。例如,对于id=1,未来 6 个月没有数据,因此 sum 将为 0。对于id=2,未来 6 个月有 1 个 num,所以total=1

我可以为每个子组执行此操作,但不能为完整的数据集执行此操作。对于一个sub-group数据框,这是我尝试过的:

 df1 = df[df['id']==1]
 
 def get_future_sum(val):
     end_date = val.date + relativedelta(months=+6)
     date_range = df1[(df1['date'] > val.date) &
                     (df1['date'] <= end_date)]
     return date_range['num'].sum()

df1['total'] = df1.apply(get_future_sum, axis=1)

最终的数据框应如下所示:

                     id               date             total
0                         1            2005-01-01           0
1                         1            2006-01-05           0
2                         2            2005-01-01           1
3                         2            2005-04-01           0
4                         3            2009-01-01           0

我尝试使用 groupby 对多个组进行迭代并应用,但它不起作用,因为df1每个组都发生了变化,我不确定如何适应它。

有没有更短的方法来遍历所有组而不为每个组使用 for 循环?

标签: pythonpandasdataframedatetime

解决方案


  • 原始问题:是否有更短的方法来遍历所有组?
    • 用于通过groupby以下方式收集数据框组id
  • 我尝试使用它在多个组上迭代它groupbyapply但它不起作用,因为df1每个组都发生了变化,我不确定如何适应它。
    • 它在通过更新发送到函数的分组数据帧正确实施时工作
  • 有没有更短的方法来遍历所有组而不为每个组使用 for 循环?
    • 可能不会,因为该函数需要每个组的数据框来确定date_range,并且该函数需要遍历每一行。
  • 遍历 groupby 对象
  • 该函数def get_future_sum对分组数据帧的每一行进行计算,并返回一个总和。
  • 将每组的结果保存在一个list
  • 用于pd.concat从数据框列表中创建单个数据框
  • 对于给定的数据,问题中显示的预期输出不正确。
    • 例如,对于id=1,没有未来 6 个月的数据,因此总和将为0
    • 对于id=2,未来 6 个月有 1 个 num,所以total=1
  • 另请参阅分组方式:split-apply-combine 用户指南
import pandas as pd
from dateutil.relativedelta import relativedelta

# test data and dataframe
data = {'id': [1, 1, 2, 2, 3],
        'date': [pd.Timestamp('2005-01-01 00:00:00'), pd.Timestamp('2006-01-05 00:00:00'), pd.Timestamp('2005-01-01 00:00:00'), pd.Timestamp('2005-04-01 00:00:00'), pd.Timestamp('2009-01-01 00:00:00')],
        'num': [7, 2, 1, 1, 2]}

df = pd.DataFrame(data)

# updated function
def get_future_sum(val: pd.Series, d: pd.DataFrame) -> np.int64:
    end_date = val.date + relativedelta(months=+6)
    date_range = d[(d['date'] > val.date) & (d['date'] <= end_date)]
    return date_range['num'].sum()


dfg = list()
for g, dg in df.groupby('id'):
    dg['total'] = dg.apply(lambda x: get_future_sum(x, dg), axis=1)
    dfg.append(dg)

# create new dataframe from list of dataframes
df_total = pd.concat(dfg).reset_index(drop=True)

# display(df_total)
   id       date  num  total
0   1 2005-01-01    7      0
1   1 2006-01-05    2      0
2   2 2005-01-01    1      1
3   2 2005-04-01    1      0
4   3 2009-01-01    2      0

推荐阅读