如何对Pandas DataFrame排序并按分类列限制每页元素数量？

问题背景 #

需要对Pandas DataFrame按数值列（如product）排序，但同时限制分页后每页中同一分类列（如seller）的元素数量。例如每页5条数据，要求每个seller在单页内最多出现2条。

原数据 #

product     seller
 10         seller1
 9          seller1
 8          seller2
 7          seller2
 6          seller2
 5          seller3

期望结果 #

product     seller
 10         seller1
 9          seller1
 8          seller2
 7          seller2
 5          seller3
 6          seller2

调整逻辑：第1页（前5条）中seller2已达到2条上限，因此将该seller的第3条记录调整至第1页之后。

解决方案 #

核心思路是先保证数值列的排序优先级，再通过标记每个分类条目出现的序号，将超出限制的条目延后排列，最终实现分页时的分类数量限制。

代码实现 #

import pandas as pd

# 构造原数据
df = pd.DataFrame({
    'product': [10, 9, 8, 7, 6, 5],
    'seller': ['seller1', 'seller1', 'seller2', 'seller2', 'seller2', 'seller3']
})

# 1. 按product降序排序，确保数值优先级
df_sorted = df.sort_values('product', ascending=False).reset_index(drop=True)

# 2. 给每个seller的条目分配组内序号（从1开始）
df_sorted['group_rank'] = df_sorted.groupby('seller').cumcount() + 1

# 3. 标记块：前2条归为块0，第3条及以后归为块1
# 若需支持更多页，可扩展为 (group_rank - 1) // 2 来生成多批次块
df_sorted['block'] = df_sorted['group_rank'].apply(lambda x: 0 if x <= 2 else 1)

# 4. 按块排序，块0在前；同块内保持product降序，最后清理辅助列
result = df_sorted.sort_values(['block', 'product'], ascending=[True, False])
result = result.drop(['group_rank', 'block'], axis=1)

print(result)

代码说明 #

1. 基础排序：先按product降序排列，确保数值列的优先级不受分类限制影响。
2. 组内序号标记：用groupby.cumcount()给每个seller的记录按出现顺序编号，区分是该seller的第几条记录。
3. 块划分：将每个seller的前2条标记为"优先块"（块0），超出的标记为"延后块"（块1），这样分页时优先块的记录会先填满页面。
4. 最终排序：按块顺序排列，同块内保持原数值排序，即可得到符合要求的结果。

扩展到多页场景 #

如果数据量更大、需要多页分页，可将块的计算逻辑调整为：

df_sorted['block'] = (df_sorted['group_rank'] - 1) // 2

这样每个seller的第1-2条在块0，第3-4条在块1，第5-6条在块2...分页时按块+每页大小分配，就能保证每页中每个seller的记录不超过2条。

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