Pandas 获取每组未来 6 个月数据的列值总和

Question

我有一个数据框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 组，我想获得未来 6 个月数据的 num 总和。比如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 循环？

Answer 1

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

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