如何对Pandas分组后的DataFrame按重复顺序进行排序？

原始数据集 #

目标排序结果 #

需求说明 #

需要让col2在每个col1分组内呈现重复序列，比如col1为'a'的分组中，col2按1、2、3、1、2、3的顺序排列，而非默认的1、1、2、2、3、3。

尝试的错误代码 #

import pandas as pd

# 创建DataFrame
data = {
    'col1': ['a', 'a', 'a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b', 'b'],
    'col2': [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6],
    'col3': ['r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', 'q', 'w', 'e']
}

df = pd.DataFrame(data)

# 先按col1、col2排序，再分组内按col2排序
df_sorted = df.sort_values(by=['col1', 'col2']).reset_index(drop=True)
df_sorted = df_sorted.groupby('col1', group_keys=False).apply(lambda x: x.sort_values('col2'))

# 输出排序后的DataFrame
print(df_sorted)

解决方案 #

原代码仅重复按col2排序，无法区分同一col2值的不同出现批次。我们可以通过给每个分组内的相同col2值添加批次序号，再按批次+col2排序来实现需求：

import pandas as pd

# 创建DataFrame
data = {
    'col1': ['a', 'a', 'a', 'a', 'a', 'a', 'b', 'b', 'b', 'b', 'b', 'b'],
    'col2': [1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6],
    'col3': ['r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z', 'q', 'w', 'e']
}

df = pd.DataFrame(data)

# 为每个(col1, col2)分组内的记录添加批次序号
df['batch'] = df.groupby(['col1', 'col2']).cumcount()

# 按col1、批次序号、col2排序，最后删除临时批次列
df_sorted = df.sort_values(by=['col1', 'batch', 'col2']).drop('batch', axis=1).reset_index(drop=True)

print(df_sorted)

代码说明 #

• groupby(['col1', 'col2']).cumcount()会在每个(col1, col2)分组内生成从0开始的递增序号，同一col2值的第一条记录批次为0，第二条为1。
• 排序时先按col1分组，再按批次序号，最后按col2，这样就能先排列所有批次0的col2序列（1、2、3），再排列批次1的col2序列（1、2、3），完全符合目标排序要求。

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