Pandas中实现兼顾成本升序与产品内数量严格升序的DataFrame排序

问题分析 #

需求核心是：每个产品下的数量必须严格按升序排列，同时整体尽可能按成本升序。这本质是一个多路归并排序问题——每个产品组内已按数量排好序，我们需要从各个组的当前待选行中，每次选出成本最小的行，直到所有行都被选取。

高效解决方案（基于堆的多路归并） #

利用Python的heapq模块实现多路归并，时间复杂度为O(n log k)（n为总行数，k为产品数量），远优于原方法的O(n²)。

import heapq
import pandas as pd

def efficient_sort(df):
    # 1. 按产品分组，每组内按数量升序排序，确保组内顺序符合要求
    grouped = df.groupby('product', sort=False).apply(
        lambda x: x.sort_values('quantity', ascending=True)
    ).reset_index(drop=True)
    
    # 2. 为每个产品创建行迭代器（按数量顺序遍历）
    product_iterators = {
        prod: iter(grouped[grouped['product'] == prod].itertuples(index=False, name=None))
        for prod in grouped['product'].unique()
    }
    
    # 3. 初始化堆：将每个产品的第一行（最小数量）放入堆，以成本为排序键
    heap = []
    for prod, it in product_iterators.items():
        try:
            row = next(it)
            heapq.heappush(heap, (row[2], prod, row))  # row[2]对应cost列
        except StopIteration:
            pass
    
    # 4. 循环弹出堆顶（成本最小）的行，同时将对应产品的下一行推入堆
    result_rows = []
    while heap:
        cost, prod, row = heapq.heappop(heap)
        result_rows.append(row)
        try:
            next_row = next(product_iterators[prod])
            heapq.heappush(heap, (next_row[2], prod, next_row))
        except StopIteration:
            pass
    
    # 5. 将结果转换为DataFrame，保持原列顺序
    return pd.DataFrame(result_rows, columns=df.columns)

# 测试示例
data = {
    'product': ['A', 'A', 'A', 'B', 'B', 'B'],
    'quantity': [0, 1, 2, 0, 1, 2],
    'cost': [0, 5, 3, 2, 4, 4]
}
df = pd.DataFrame(data)
result = efficient_sort(df)
print(result)

代码解释 #

1. 分组排序：先确保每个产品组内的行按quantity升序排列，满足产品内数量严格递增的要求。
2. 迭代器初始化：为每个产品创建迭代器，保证按数量顺序逐行取出。
3. 堆初始化：将每个产品的第一行（最小数量）放入堆，堆自动按cost升序维护，确保每次能快速取出当前成本最低的行。
4. 堆循环处理：每次弹出堆顶行加入结果，再将该产品的下一行（如果存在）推入堆，直到所有行都被处理。

效率对比 #

• 原方法：逐行遍历并搜索idxmin，时间复杂度O(n²)，当数据量较大（如万级以上行）时，性能会急剧下降。
• 新方法：基于堆的多路归并，时间复杂度O(n log k)，k为产品数量，即使处理百万级行数据，也能保持高效。

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