pandas如何按列值优先级从DataFrame分组选取指定数量行

核心思路 #

全程使用向量化操作保证运行性能，无逐行循环逻辑，实现步骤如下：

• 按照给定的val1、val2组合规则，为每一行计算优先级权重，优先级越高权重值越小
• 按ID字段分组，组内按优先级权重升序排列
• 每个分组截取最多2行，删除临时生成的优先级列即可得到最终结果

完整实现代码 #

import pandas as pd
import numpy as np

# 构造原始示例数据
df = pd.DataFrame({
    'ID': ['A', 'A', 'B', 'B', 'B', 'c', 'c'],
    'val1': [10, 10, 10, 9, 8, 9, 9],
    'val2': [0, 0, 0, 1, 2, 1, 1],
    'uniq': ['1-2', '3-2', '8-0', '7-6', '9-3', '10-3', '3-0']
})

# 配置优先级规则：条件越靠前优先级越高，对应权重值越小
priority_conds = [
    (df['val1'] == 10) & (df['val2'] == 0),
    (df['val1'] == 9) & (df['val2'] == 1),
    (df['val1'] == 8) & (df['val2'] == 2)
]
priority_weights = [1, 2, 3]
# 不满足任何优先级规则的行默认给极低优先级（权重999）
df['p_rank'] = np.select(priority_conds, priority_weights, default=999)

# 分组排序后每个ID取前2行，删除临时优先级列
result = df.sort_values(by=['ID', 'p_rank']).groupby('ID', as_index=False).head(2).drop(columns='p_rank')

输出验证 #

执行代码后result的输出和期望结果完全一致：

ID  val1  val2 uniq
0  A    10     0  1-2
1  A    10     0  3-2
2  B    10     0  8-0
3  B     9     1  7-6
5  c     9     1 10-3
6  c     9     1  3-0

补充说明 #

• 该实现使用numpy、pandas原生向量化接口，十万级以上数据量下运行效率远高于apply类逐行处理方案
• 如果同一优先级下需要按其他字段（比如uniq列）排序决定取哪些行，直接在sort_values的by参数里追加对应字段即可
• 如果后续要调整优先级规则，只需要修改priority_conds和priority_weights两个列表的对应项，不需要改动核心逻辑

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