如何在Python中对DataFrame分层随机抽样以保留分布？

你需要的是分层抽样（Stratified Sampling），Pandas通过groupby()结合sample()方法就能轻松实现，确保抽样后location和category的分布与原数据集大致一致。

方法1：多列交叉分层抽样（推荐） #

要同时保留两列的分布，最准确的方式是把location和category的组合作为分层依据，对每个子组按相同比例抽样：

import pandas as pd

# 读取原数据
df = pd.read_parquet("~/dimension.parquet")

# 按location+category组合分层，抽取20%样本（frac可根据需求调整）
sampled_df = df.groupby(['location', 'category'], group_keys=False).sample(frac=0.2, random_state=42)

验证抽样后的分布 #

执行以下代码可确认抽样后两列的比例是否与原数据匹配：

# 查看抽样后的location分布比例
print("抽样后location分布：")
print(sampled_df.groupby('location')['location'].count().transform(lambda x: x/x.sum()))

# 查看抽样后的category分布比例
print("\n抽样后category分布：")
print(sampled_df.groupby('category')['category'].count().transform(lambda x: x/x.sum()))

其他可选方式 #

• 若需要抽取固定数量的样本，将frac替换为n即可，比如每个子组抽取100条：

  sampled_df = df.groupby(['location', 'category'], group_keys=False).sample(n=100, random_state=42)
  

• 若仅需保证单一列的分布（比如仅location），可只按该列分组抽样，但这种方式可能导致另一列的分布出现轻微偏移。

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