保留年龄/性别分布裁剪Pandas DataFrame至约10k条数据

不需要用循环遍历每个分组，Pandas本身就有更简洁高效的分层抽样方法，能精准保留age和sex的联合分布，同时快速缩减到目标样本量：

方法一：分组比例抽样（兼容所有Pandas版本） #

先计算整体抽样比例，再按age+sex的组合分组，每组按比例抽取样本，最后合并结果。这种方法能严格维持原数据的分组分布比例：

target_size = 10000
# 计算抽样比例
sample_ratio = target_size / len(base)

# 按年龄和性别分组抽样，group_keys=False避免保留分组索引
sampled_df = base.groupby(['age', 'sex'], group_keys=False).apply(
    lambda group: group.sample(frac=sample_ratio, random_state=42)
)

# 若最终样本量与目标有微小偏差，可根据需求调整（比如直接取前10000条）
# sampled_df = sampled_df.head(target_size)

方法二：直接使用strata参数（Pandas 1.3.0+） #

Pandas 1.3.0及以上版本的sample方法支持strata参数，可直接指定分层列，代码更简洁：

target_size = 10000
sampled_df = base.sample(n=target_size, strata=['age', 'sex'], random_state=42)

注意事项 #

• random_state参数用于固定抽样结果，保证可复现，不需要的话可以删除；
• 如果部分age+sex分组的样本量极小，抽样时可以添加replace=True允许放回抽样（但非必要不推荐，会引入重复数据）；
• 两种方法都比循环实现效率高几个数量级，尤其适合大规模数据集。

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