如何扩展pandas DataFrame使分组内每行包含此前所有行

需求说明 #

对DataFrame中的每个分组实现行的累积追加，效果与下文手动实现的示例一致，约束条件为不使用for循环或df.apply()方法，仅通过向量化操作单次实现。

示例代码与预期输出 #

原始输入df1：

import pandas as pd
import numpy as np

df1 = pd.DataFrame(np.array([[1, 1, 1], [2, 2, 2], [3, 3, 3]]),
                   columns=['group', 'a', 'b'])
df1 = df1.set_index('group').sort_index()
print(df1)

a  b
group      
1      1  1
2      2  2
3      3  3

预期输出df2：

df2 = pd.DataFrame(np.array([[1, 1, 1], [2, 1, 1], [2, 2, 2], [3, 1, 1], [3, 1, 1], [3, 2, 2]]),
                   columns=['group', 'a', 'b'])
df2 = df2.set_index('group').sort_index()
print(df2)

a  b
group      
1      1  1
2      1  1
2      2  2
3      1  1
3      1  1
3      2  2

实现方案 #

核心逻辑是通过numpy向量化重复生成基础索引框架，给分组内数据打序号后通过关联匹配+前向填充实现累积追加效果，全程无逐组循环、无df.apply()调用，执行效率为向量化运算级别。

# 重置索引取回group字段
df1 = df1.reset_index()

# 1. 统计每个分组需要生成的总行数，示例规则为group值为g则生成g行，通用场景下直接取组内原始行数即可
group_cnt = df1.groupby('group', as_index=False).size()
group_cnt['size'] = group_cnt['group'] # 若为通用累积场景，注释掉这行即可，默认按组内原始行数生成对应长度的累积序列

# 2. 生成包含分组内序号的基础表
base = pd.DataFrame({
    'group': np.repeat(group_cnt['group'], group_cnt['size'])
})
base['seq'] = base.groupby('group').cumcount() + 1

# 3. 给原始表的分组内行打序号
df1['seq'] = df1.groupby('group').cumcount() + 1

# 4. 关联匹配+前向填充得到累积追加结果
result = base.merge(df1, on=['group', 'seq'], how='left')\
             .ffill()\
             .drop(columns='seq')\
             .set_index('group')\
             .sort_index()

执行后打印result即可得到和示例完全一致的df2效果。

通用场景适配：如果你的原始数据每个分组有n行，需要实现"第1行保留1行，前2行追加为2行，...，前n行追加为n行"的全量累积效果（单组最终行数为n(n+1)/2），只需要注释掉group_cnt['size'] = group_cnt['group']这一行，其余代码无需修改即可直接使用。

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