如何基于groupby键动态设置sample的n参数从分组DataFrame抽取随机样本

需求说明 #

从groupby后的DataFrame对象中抽取随机样本，需根据分组键（cc）对应的selected_count值动态设置sample()函数的n参数：

• US分组抽取3行（对应selected_count=3）
• UK分组抽取1行（对应selected_count=1）

示例数据 #

创建DataFrame代码 #

import pandas as pd

d = {'name': ["n1", "n2", "n3", "n4", "n5", "n6"], 
     'cc': ["US", "UK", "US", "UK", "US", "US"], 
     'selected_count':[3, 1, 3, 1, 3, 3], 
     'view':[4, 64, 52, 2, 65, 21]}
pdf_candidate_names = pd.DataFrame(data=d)

数据框输出 #

name  cc  selected_count  view
0   n1  US               3     4
1   n2  UK               1    64
2   n3  US               3    52
3   n4  UK               1     2
4   n5  US               3    65
5   n6  US               3    21

失败尝试及原因 #

失败代码 #

pdf_selected_names = pd.concat([
    pdf_candidate_names.groupby("cc").apply(lambda x: x.sample(n=x["selection_count"], weights='views')),
    pdf_candidate_names.groupby("cc").apply(lambda x: x.sample(n=x["selection_count"], weights='views'))
]).sample(frac=1.0).reset_index(drop=True)

失败原因 #

1. 列名错误：代码中x["selection_count"]应为x["selected_count"]，weights='views'应为weights='view'（与数据列名一致）
2. 核心问题：x["selected_count"]返回的是Series对象，而sample()的n参数需要单个整数，无法直接传入。

解决方案 #

由于每个分组内的selected_count值是统一的（如US组全为3，UK组全为1），只需取出分组内该列的任意一个值作为整数传入n即可：

正确代码 #

import pandas as pd

# 先确保每个分组的selected_count值统一（可选验证步骤）
assert pdf_candidate_names.groupby("cc")["selected_count"].nunique().eq(1).all(), "部分分组的selected_count值不统一"

# 按分组动态抽样
pdf_selected_names = pdf_candidate_names.groupby("cc", group_keys=False).apply(
    lambda x: x.sample(n=x["selected_count"].iloc[0], weights='view')
).reset_index(drop=True)

代码说明 #

• group_keys=False：避免分组键cc被添加为结果的索引层级
• x["selected_count"].iloc[0]：取出当前分组内selected_count列的第一个值（因同组值统一，取任意位置均可），得到整数型的样本量
• weights='view'：按view列的值作为抽样权重（可选，若不需要加权可删除该参数）

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