GroupBy后如何对DataFrame按分组进行子抽样处理？

现有如下结构的DataFrame：

import pandas as pd
df = pd.DataFrame({
    'id': [1,1,1,1,2,2,2,3,3,4,4],
    'tweet': ['a','b','b','a','d','a','a','b','b','a','b']
})

按id分组统计推文数量：

df.groupby(["id"]).count()

得到结果：

需求：对子抽样，保留推文数小于n的用户的全部数据，对推文数超过n的用户随机抽取n条推文。尝试以下代码后发现仅返回整个数据集的n条样本，求正确实现方式：

n=3
print(data.groupby(["user_id"]).apply(lambda x: x.sample(min(n,len(x)), replace=False)).reset_index(drop=True))
print(data.groupby('user_id').sample(n, random_state=1))

你写的代码有两个关键问题：

1. 变量名不匹配：原数据框是df，代码里误用了data；原分组列是id，代码里写成了user_id；
2. groupby.sample参数设置错误：直接传n的话，当分组长度小于n时会报错，且无法实现“短分组全保留、长分组抽n条”的逻辑。

下面是两种可行的正确代码：

方法1：结合groupby.apply和sample #

n = 3
# 按id分组，每个分组取min(n, 分组长度)条数据
result = df.groupby("id").apply(lambda x: x.sample(min(n, len(x)), replace=False)).reset_index(drop=True)
print(result)

方法2：用groupby.sample的lambda参数（Pandas 1.1.0及以上版本支持） #

n = 3
# 给n传一个lambda函数，动态计算每个分组要抽取的数量
result = df.groupby("id").sample(n=lambda x: min(n, len(x)), random_state=1).reset_index(drop=True)
print(result)

运行后会得到符合需求的子抽样数据：

• id=1（原4条）：随机抽取3条；
• id=2（原3条）：保留全部3条；
• id=3、id=4（各2条）：保留全部数据；
• 最终结果总条数为10条（id=1的3条是随机的，每次运行结果可能略有不同）。

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