基于列递增规律分组Pandas DataFrame并生成累积列表列

原始数据 #

import pandas as pd

df = pd.DataFrame({
    'col1': [0, 0, 0, 1, 0, 1, 0, 0, 1],
    'col2': [1, 2, 3, 4, 1, 2, 1, 2, 3]
})

需求说明 #

根据col2的递增规律分组（当col2值小于前一行时开启新组），新增col3列，存储对应行及之前同组内col2的累积数值列表。

期望结果 #

result = pd.DataFrame({
    'col1': [0, 0, 0, 1, 0, 1, 0, 0, 1],
    'col2': [1, 2, 3, 4, 1, 2, 1, 2, 3],
    'col3': [[1], [1,2], [1,2,3], [1,2,3,4], [1], [1,2], [1], [1,2], [1,2,3]]
})

简洁高效实现方法 #

# 生成分组标识：当col2不大于前一行时标记为新组，累积求和得到组ID
df['group'] = df['col2'].diff().le(0).cumsum()

# 对每个组生成col2的累积列表
df['col3'] = df.groupby('group')['col2'].expanding().apply(list, raw=False).reset_index(level=0, drop=True)

# 移除临时分组列（可选）
df = df.drop('group', axis=1)

代码解释 #

• 分组键生成：df['col2'].diff().le(0)判断当前col2是否不大于前一行值，返回布尔序列；cumsum()将布尔值转为1/0累加，得到唯一的组ID，每个连续递增的col2序列对应一个组。
• 累积列表生成：通过groupby按组拆分后，用expanding()窗口覆盖从组内第一行到当前行的所有数据，apply(list)将窗口内的col2值转为列表；最后重置索引把结果合并回原DataFrame。

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