如何将col_A分组的col_B最大值分配给下一组所有行？

问题描述 #

现有如下DataFrame：

需求：将col_A中某分组的col_B最大值，分配给col_A升序排列的下一组的所有行，期望输出如下：

用户原始思路为：先按分组获取max值，再将结果shift后合并回源DataFrame，但不清楚具体实现方式。要求方案需适配150万行数据、8万个col_A离散值的高效处理。

DataFrame复现代码：

import pandas as pd

df = pd.DataFrame(
    {
        'col_A': {0: 1, 1: 1, 2: 1, 3: 2, 4: 2, 5: 2, 6: 3, 7: 3, 8: 3},
        'col_B': {0: 1, 1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: 7, 7: 8, 8: 9},
    }
)

高效解决方案 #

实现步骤 #

1. 按col_A分组计算col_B的最大值，得到每个分组对应的最大值Series
2. 对分组最大值执行shift(1)操作，将上一组的最大值映射到当前组
3. 通过map方法将处理后的最大值关联回原DataFrame，生成目标列

代码实现 #

# 计算每个col_A分组的col_B最大值
group_max = df.groupby('col_A')['col_B'].max()

# 移位操作：将上一组的最大值对应到当前组
shifted_max = group_max.shift(1)

# 映射回原DataFrame，生成目标列
df['max_col_B_from_prev_col_A_group'] = df['col_A'].map(shifted_max)

方案优势 #

• 高效聚合：groupby.max()是Pandas底层优化的聚合操作，处理百万级数据性能优异
• 轻量移位：shift()仅对Series进行索引偏移，时间复杂度为O(n)，无额外开销
• 快速映射：map()基于哈希表实现查找，针对8万个离散值的匹配操作远快于表连接（merge），避免大规模数据的连接开销

注意事项 #

如果原DataFrame的col_A未按升序排列，需先执行排序确保分组顺序正确：

# 对col_A升序排序（根据实际数据情况选择是否执行）
df = df.sort_values('col_A').reset_index(drop=True)

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