Pandas按指定列表对DataFrame分组排序并保留重复值

需求说明 #

现有如下结构的DataFrame：

需要基于指定排序列表levels = [1,2,3]，对每个org分组内的数据排序，保留所有重复值，最终得到分组内按1→2→3循环排列的结果：

实现方案 #

核心逻辑是给每个org分组下、相同level的重复值标记出现轮次，先按轮次排序，再按指定的level顺序排序，即可实现循环排列的效果，不会丢失任何行。

完整可运行代码 #

import pandas as pd

# 构造测试数据
df = pd.DataFrame({
    'level': [1,1,2,2,3,3]*3,
    'org': ['A']*6 + ['B']*6 + ['C']*6
})
levels = [1,2,3]

# 步骤1：给同org下同level的行标记出现轮次（辅助列）
df['round'] = df.groupby(['org', 'level']).cumcount()

# 步骤2：按规则排序，删除辅助列后重置索引
df_result = (
    df
    # 先按org分组，再按轮次排序，最后按自定义level顺序排序
    .sort_values(
        by=['org', 'round', 'level'],
        # 对level列应用自定义排序顺序，其余列保持默认排序规则
        key=lambda x: x.map({val: idx for idx, val in enumerate(levels)}) if x.name == 'level' else x
    )
    .drop(columns='round')
    .reset_index(drop=True)
)

逻辑拆解 #

• 分组累计计数：groupby(['org', 'level']).cumcount()会给每个org下、每个level的重复值从0开始编号，例如org=A下两个level=1的行编号为0、1，两个level=2的行编号也为0、1，以此类推。
• 多键排序：排序时优先按org划分分组，同一org下先按轮次round排序，保证第一轮的所有level排完再排第二轮；同一轮次下按指定的levels顺序排列level值，自然得到1→2→3的循环顺序。
• 清理辅助字段：排序完成后删除用于标记轮次的临时列，重置索引即可得到目标结果。

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