基于Python实现覆盖全量国家与州的DataFrame随机抽样需求

核心思路 #

要实现覆盖所有唯一Country-State组合、每个组合至少包含一个bill_ID，同时总抽样量为5000行的需求，分两步执行：

• 先提取每个Country-State组至少一个完整bill_ID的所有关联行，确保覆盖要求；
• 从剩余数据中随机抽取补足至5000行。

代码实现（Pandas） #

import pandas as pd
import numpy as np

# 假设原始数据集为df
# 1. 按Country-State分组，每组随机选一个bill_ID
grouped = df.groupby(['Country', 'State'])
selected_bills = grouped['bill_ID'].apply(lambda x: np.random.choice(x.unique())).reset_index()

# 获取选中bill_ID对应的所有行（保证每个Country-State至少有一个完整bill的记录）
mandatory_rows = df[df['bill_ID'].isin(selected_bills['bill_ID'])]

# 2. 计算需要额外抽取的行数
remaining_needed = 5000 - len(mandatory_rows)

# 3. 补足剩余抽样量
if remaining_needed > 0:
    # 筛选出不在必选bill_ID中的行
    candidate_rows = df[~df['bill_ID'].isin(selected_bills['bill_ID'])]
    # 随机抽取剩余行数（random_state保证结果可复现，适合测试）
    sampled_extra = candidate_rows.sample(n=remaining_needed, random_state=42)
    # 合并得到最终样本
    final_sample = pd.concat([mandatory_rows, sampled_extra])
else:
    # 极端情况：必选行已超过5000，直接从必选行中随机抽5000
    final_sample = mandatory_rows.sample(n=5000, random_state=42)

# 重置索引
final_sample = final_sample.reset_index(drop=True)

关键细节说明 #

• 用groupby+np.random.choice确保每个Country-State组都有至少一个随机选中的bill_ID，且取出该bill_ID的所有关联行（符合一个bill_ID对应多个item_id的业务逻辑）；
• 设置random_state=42是为了让抽样结果可复现，方便测试场景下的重复验证；
• 处理了必选行总数超过5000的极端情况，保证最终样本量严格符合要求。

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