Pandas按A列分层抽样且保证B列所有唯一值至少出现一次的实现方法

前置说明 #

你现有代码存在两个基础问题：一是循环索引使用i-1取值，当i从0开始时会错误取到列表最后一个元素；二是未考虑B列唯一值的保底覆盖需求，纯随机抽样必然会出现小频率B值遗漏。
另外首先要确认：你的总抽样量desired_sample_size必须 ≥ B列唯一值数量（你的业务场景下至少要≥22600），否则不可能满足所有B唯一值都出现一次的硬约束。

核心实现逻辑 #

采用「保底优先、补量对齐占比」的两阶段抽样：

• 第一阶段：对每个B的唯一值，随机抽取1条对应记录进入样本池，确保B列100%覆盖
• 第二阶段：
  1. 统计保底阶段各A层已抽取的样本量
  2. 按照预设的A层占比计算各层最终目标样本量
  3. 计算各层还需补抽的样本量（若某层保底量已超过目标量，可微调其他层配额，或小幅提升总抽样量）
  4. 从原数据中排除已被抽取的保底样本，按层补抽对应数量的样本，合并后得到最终样本

完整可运行代码 #

import pandas as pd
import random
import string
import numpy as np

# ---------------------- 模拟数据生成 ----------------------
strata_groups=['st1', 'st2', 'st3', 'st4']
A = random.choices(strata_groups, k=10000)
B = random.choices(string.ascii_lowercase,k=10000)
C = np.random.randint(1, 6, 10000)
original_data=pd.DataFrame({'A': A, 'B': B, 'C': C})

# ---------------------- 抽样参数配置 ----------------------
strata_names_under_col_A=['st1', 'st2', 'st3', 'st4']
strata_prop=[0.25, 0.40, 0.15, 0.20] # A列分层抽样预设占比
desired_sample_size=1000 # 总抽样量，必须≥B列唯一值数量，此处模拟数据B有26个唯一值，1000满足要求
random_seed=42 # 固定随机种子保证可复现

# ---------------------- 第一阶段：B列保底抽样 ----------------------
# 每个B唯一值随机抽1条
base_sample = original_data.groupby('B', group_keys=False).apply(lambda x: x.sample(1, random_state=random_seed))
used_index = set(base_sample.index) # 记录已被抽取的索引，避免重复抽样

# ---------------------- 第二阶段：补量抽样对齐A层占比 ----------------------
# 计算各A层目标样本量
strata_target = {name: round(desired_sample_size * prop) for name, prop in zip(strata_names_under_col_A, strata_prop)}
# 修正目标总量误差（四舍五入可能导致总和不等于desired_sample_size）
diff = desired_sample_size - sum(strata_target.values())
if diff != 0:
    strata_target[strata_names_under_col_A[0]] += diff

# 统计保底样本中各A层的已有数量
strata_base_count = base_sample['A'].value_counts().to_dict()
# 初始化补抽样本容器
add_sample = pd.DataFrame(columns=original_data.columns)

# 按层补抽
for strata_name in strata_names_under_col_A:
    # 计算该层需要补抽的数量
    need = strata_target[strata_name] - strata_base_count.get(strata_name, 0)
    if need <= 0:
        continue # 该层保底量已足够，无需补抽
    # 筛选该层未被抽取的记录
    strata_data = original_data[(original_data['A'] == strata_name) & (~original_data.index.isin(used_index))]
    # 补抽，若剩余数据不够则允许有放回抽样（可根据需求调整）
    sample = strata_data.sample(n=need, replace=len(strata_data) < need, random_state=random_seed)
    add_sample = pd.concat([add_sample, sample], ignore_index=True)
    used_index.update(sample.index)

# 合并得到最终样本
final_sample = pd.concat([base_sample, add_sample], ignore_index=True)

# ---------------------- 结果校验 ----------------------
print(f"B列唯一值覆盖率：{len(final_sample['B'].unique()) / len(original_data['B'].unique()) * 100:.2f}%")
print("各A层实际占比：")
print(final_sample['A'].value_counts(normalize=True).reindex(strata_names_under_col_A).round(4))
print(f"总样本量：{len(final_sample)}")

适配业务场景的说明 #

• 你业务场景下B有22600个唯一值，建议desired_sample_size设置为23000及以上，预留足够的补量空间来对齐A层占比
• 若部分B值仅存在于某一个A层，保底抽样会自动占用该层的配额，补抽阶段会自动调整其他层的补量，尽可能贴合预设占比
• 如果对A层占比的精度要求极高，可以在补抽完成后，对超量的A层随机删除多余样本（优先删除重复B值的记录），进一步对齐占比

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