如何在groupby后将后续组首行拼接至当前组每一行？

需求说明 #

将每个ID分组之后出现的首个分组的首行数据，拼接至当前分组的每一行中；最后一个分组的对应Next字段填充为NaN。

输入数据 #

ID     value      counts      date
1      1          3           1/2/2020
1      2          10          10/2/2020
1      3          5           15/2/2020
2      1          6           3/4/2020
2      2          2           10/4/2020

解决方案（基于Pandas） #

1. 提取每个分组的首行数据
   按ID分组后取出每个组的第一行，保留原列名用于后续关联：

   import pandas as pd
   
   # 构造输入DataFrame
   df = pd.DataFrame({
       'ID': [1,1,1,2,2],
       'value': [1,2,3,1,2],
       'counts': [3,10,5,6,2],
       'date': ['1/2/2020', '10/2/2020', '15/2/2020', '3/4/2020', '10/4/2020']
   })
   
   # 获取每个ID组的首行数据
   group_headers = df.groupby('ID').first().reset_index()
   

2. 生成下一组的关联信息
   对分组首行数据使用shift(-1)获取下一组的首行内容，并重命名列名区分：

   # 生成下一组信息并修改列名
   next_group = group_headers.shift(-1).rename(
       columns={
           'ID': 'NextID',
           'value': 'NextValue',
           'counts': 'NextCounts',
           'date': 'NextDate'
       }
   )
   
   # 合并原组首行与下组信息
   group_with_next = pd.concat([group_headers, next_group], axis=1)
   

3. 将下组信息合并到原数据
   通过ID字段将原数据与关联好的分组信息合并，实现当前组每一行都带上下组首行数据：

   # 合并到原DataFrame
   result = df.merge(group_with_next[['ID', 'NextID', 'NextValue', 'NextCounts', 'NextDate']], on='ID', how='left')
   

输出结果 #

ID  value  counts        date  NextID  NextValue  NextCounts    NextDate
1      1       3    1/2/2020     2.0        1.0         6.0    3/4/2020
1      2      10   10/2/2020     2.0        1.0         6.0    3/4/2020
1      3       5   15/2/2020     2.0        1.0         6.0    3/4/2020
2      1       6    3/4/2020     NaN        NaN         NaN         NaN
2      2       2   10/4/2020     NaN        NaN         NaN         NaN

注：若需要包含示例中的ID=3分组，只需先将该组数据追加到原df中，再执行上述流程即可得到完整的目标输出。

内容的提问来源于stack exchange，提问作者Achille G