基于Pandas按A列连续相同值分组计算B列最大值的需求

原始数据 #

import pandas as pd

df = pd.DataFrame({
    'day': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    'A': [1, 1, 1, 0, 0, 0, 1, 1, 0, 0],
    'B': [7, 8, 4, 3, 5, 6, 3, 1, 4, 5],
})

需求说明 #

将A列中连续的相同值（连续1或连续0）作为分组，计算每组内B列的最大值，并将该最大值填充到对应分组的每一行，最终得到包含Output列的结果。

期望输出 #

expected_df = pd.DataFrame({
    'day': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
    'A': [1, 1, 1, 0, 0, 0, 1, 1, 0, 0],
    'B': [7, 8, 4, 3, 5, 6, 3, 1, 4, 5],
    'Output': [8, 8, 8, 6, 6, 6, 3, 3, 5, 5]
})

解决方案 #

核心思路是先生成连续相同值的分组标识，再通过分组计算最大值并填充到对应行：

# 生成连续分组的唯一标识：A列值变化时分组键递增
group_key = df['A'].ne(df['A'].shift()).cumsum()

# 按分组键计算每组B的最大值，并用transform将最大值填充到每组所有行
df['Output'] = df.groupby(group_key)['B'].transform('max')

# 查看结果
print(df)

代码说明 #

• df['A'].ne(df['A'].shift())：生成布尔序列，当前行A值与上一行不同时为True
• .cumsum()：将布尔序列转换为递增整数，每个连续相同值组对应唯一的整数键
• groupby(group_key)['B'].transform('max')：按分组键计算每组B的最大值，transform方法会将最大值映射回原分组的每一行，直接生成Output列

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