Pandas中结合Groupby与Rolling生成矩阵的实现需求

我是Pandas新手，这个问题卡了两天，试了各种方法查资料都没解决。现有如下DataFrame：

import pandas as pd

df = pd.DataFrame({'a': [1,1,1,2,2,2,3,3,3,1,1,1,2,2,2,3,3,3],
                   'b': ['b1','b1','b1','b1','b1','b1','b1','b1','b1','b2','b2','b2','b2','b2','b2','b2','b2','b2',],
                   'c': ['c1','c2','c3','c1','c2','c3','c1','c2','c3','c1','c2','c3','c1','c2','c3','c1','c2','c3',],
                   'value': ['v1','v2','v3','v4','v5','v6','v7','v8','v9','v10','v11','v12','v13','v14','v15','v16','v17','v18']})
                   
print(df)

打印后的原始DataFrame：

a   b   c value
0   1  b1  c1    v1
1   1  b1  c2    v2
2   1  b1  c3    v3
3   2  b1  c1    v4
4   2  b1  c2    v5
5   2  b1  c3    v6
6   3  b1  c1    v7
7   3  b1  c2    v8
8   3  b1  c3    v9
9   1  b2  c1   v10
10  1  b2  c2   v11
11  1  b2  c3   v12
12  2  b2  c1   v13
13  2  b2  c2   v14
14  2  b2  c3   v15
15  3  b2  c1   v16
16  3  b2  c2   v17
17  3  b2  c3   v18

需求说明 #

按b字段分组，每个分组内对连续2个a进行滚动操作，生成以a（格式为a1/a2/a3）为行、c为列、value为单元格值的矩阵，预期输出如下：

c1 c2 c3
b1   a1   v1 v2 v3
     a2   v4 v5 v6 

b1   a2   v4 v5 v6
     a3   v7 v8 v9

b2   a1   v10 v11 v12
     a2   v13 v14 v15

b2   a2   v13 v14 v15
     a3   v16 v17 v18

解决方案 #

实现步骤 #

• 先对原始数据做透视，将每个(b,a)组合对应的c列值整合为一行，方便后续滚动操作
• 按b分组后，对每个分组内的行执行窗口大小为2的滚动遍历
• 格式化输出每个滚动窗口对应的矩阵

代码实现 #

import pandas as pd

# 原始数据
df = pd.DataFrame({'a': [1,1,1,2,2,2,3,3,3,1,1,1,2,2,2,3,3,3],
                   'b': ['b1','b1','b1','b1','b1','b1','b1','b1','b1','b2','b2','b2','b2','b2','b2','b2','b2','b2',],
                   'c': ['c1','c2','c3','c1','c2','c3','c1','c2','c3','c1','c2','c3','c1','c2','c3','c1','c2','c3',],
                   'value': ['v1','v2','v3','v4','v5','v6','v7','v8','v9','v10','v11','v12','v13','v14','v15','v16','v17','v18']})

# 1. 透视数据：将(b,a)作为索引，c作为列，value作为值
pivot_df = df.pivot(index=['b', 'a'], columns='c', values='value').reset_index()
# 把a列格式化为a1、a2的形式
pivot_df['a'] = pivot_df['a'].apply(lambda x: f'a{x}')

# 2. 按b分组处理滚动窗口
for b_group, sub_df in pivot_df.groupby('b'):
    # 遍历每个连续2行的窗口
    for i in range(len(sub_df) - 1):
        # 提取当前窗口的行，设置a为索引并移除b列
        window_data = sub_df.iloc[i:i+2].set_index('a').drop('b', axis=1)
        # 输出分组标识和窗口矩阵
        print(f"{b_group}")
        print(window_data)
        print()

代码说明 #

• pivot操作将分散的c列值整合到同一行，解决了原始数据中每个(b,a,c)对应单个value的问题
• 按b分组后，通过遍历索引实现窗口大小为2的滚动，确保每个分组内只取连续的两个a组
• 格式化a列为a1/a2样式，与预期输出一致

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