在Pandas中基于已排序列实现低复杂度双列排序

问题背景 #

现有大型Pandas DataFrame，无法承受全量O(n²)复杂度的排序操作。已知col1列已处于有序状态，需实现按col1和col2两列排序的逻辑，要求最终结果与df.sort_values(['col1', 'col2'])完全一致。

由于仅需对col1的分组内数据按col2排序，且每组元素数量k范围为1-12，整体时间复杂度可降至O(n(klogk))，远优于全量排序的最坏O(n²)复杂度。

效果验证示例 #

以下是测试用例及结果验证逻辑：

import pandas as pd

df = pd.DataFrame({
    'col1': [0, 0, 1, 2, 2, 2],
    'col2': ['c', 'a', 'b', 'e', 'h', 'f']
})

# 调用自定义排序函数
df2 = sort_ordered(df, 'col1', 'col2')

# 验证结果与原生排序完全一致
assert df2.equals(df.sort_values(['col1', 'col2']))

迭代实现思路（修正版） #

核心逻辑是遍历DataFrame，定位col1分组的起止索引，对每个分组内的col2单独排序。原伪代码存在逻辑瑕疵，以下是修正后的实现思路伪代码：

wstart = 0
# 遍历到最后一行之后，处理收尾的最后一个分组
for i in range(len(df) + 1):
    if i == len(df) or df.iloc[i]['col1'] != df.iloc[wstart]['col1']:
        # 对当前分组（wstart到i-1区间）按col2排序
        df.iloc[wstart:i] = df.iloc[wstart:i].sort_values('col2')
        wstart = i

思路说明 #

• 用wstart标记当前分组的起始索引
• 遍历过程中，当遇到col1值变化或到达DataFrame末尾时，对当前分组执行col2排序
• 排序完成后更新wstart为下一个分组的起始索引

内容的提问来源于stack exchange，提问作者David Davó