Pandas DataFrame条件分组筛选：按规则选取组内指定数量子集

我来帮你实现这个需求！核心思路是先锁定必须保留的include=1行，再用剩余名额从include=0里挑选Rank最小的行，确保每组最多3行，同时总Rank最小。下面是完整的代码和步骤解释：

步骤1：构造示例DataFrame #

首先咱们先把你给出的数据转换成Pandas DataFrame：

import pandas as pd

df = pd.DataFrame({
    'ID': [1,2,3,4,5,6,7,8,9],
    'Group': ['Group1','Group1','Group1','Group1','Group2','Group2','Group2','Group2','Group2'],
    'Rank': [1,2,3,4,1,2,3,4,5],
    'include': [0,1,0,1,0,0,1,1,1]
})

步骤2：拆分必须保留和候选数据 #

先把include=1的行单独拎出来（这部分是必须全部保留的），剩下的include=0作为候选补充行：

# 必须保留的行：所有include=1的记录
required_rows = df[df['include'] == 1].copy()
# 候选补充行：include=0的记录，用于填补剩余名额
candidate_rows = df[df['include'] == 0].copy()

步骤3：计算每组剩余名额 #

每组最多选3行，减去已经保留的include=1行数，得到还能选多少候选行：

# 统计每组已保留的include=1行数
count_required = required_rows.groupby('Group').size()
# 计算每组剩余名额：3 - 已保留行数，只保留名额>0的组
remaining_quota = 3 - count_required
remaining_quota = remaining_quota[remaining_quota > 0]

步骤4：从候选行中挑选最小Rank的记录 #

对候选行按组分组，每组取前N个Rank最小的行（N是剩余名额）：

# 按组筛选候选行，每组取剩余名额数量的最小Rank记录
selected_candidates = candidate_rows.groupby('Group').apply(
    lambda group: group.nsmallest(remaining_quota.get(group.name, 0), 'Rank')
).reset_index(drop=True)

这里用remaining_quota.get(group.name, 0)是为了兼容那些已经满3个include=1行的组（比如Group2），直接返回0，不会额外选候选行。

步骤5：合并结果并排序 #

把必须保留的行和选中的候选行合并，再按组和Rank排序，得到最终结果：

# 合并两部分数据
final_result = pd.concat([required_rows, selected_candidates])
# 按Group和Rank排序，让结果更清晰
final_result = final_result.sort_values(['Group', 'Rank']).reset_index(drop=True)

最终输出 #

运行后final_result的结果如下：

ID    Group  Rank  include
0   1  Group1     1        0
1   2  Group1     2        1
2   4  Group1     4        1
3   7  Group2     3        1
4   8  Group2     4        1
5   9  Group2     5        1

结果验证： #

• Group1：include=1有2行（ID2、4），剩余1个名额，选了Rank最小的ID1，总Rank=1+2+4=7，是最小可能的组合
• Group2：include=1刚好3行，所以没有选候选行，完全符合要求

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