如何用Python对数据集进行带双布尔条件约束的随机采样？

核心思路 #

先将数据集按cond_1和cond_2的组合拆分为4个互斥分组，再根据约束条件推导每个分组的抽取数量，最后从各分组抽取对应样本并合并。

分组定义 #

• 组A：cond_1=True且cond_2=True
• 组B：cond_1=True且cond_2=False
• 组C：cond_1=False且cond_2=True
• 组D：cond_1=False且cond_2=False

约束推导 #

根据你的采样要求，可推导出各分组抽取数量的关系：

• 组A抽取数 a = 181 - b（b为组B的抽取数）
• 组C抽取数 c = b（因cond_2=True总数为181，即a + c = 181）
• 组D抽取数 d = 485 - b（因cond_1=False总数为485，即c + d = 485）
• 为满足「包含交叉组合」要求，需保证b ≥ 1（避免仅存在组A和组D的同状态样本）

同时需确保各分组抽取数不超过原数据对应分组的样本量，且a ≥ 0、d ≥ 0，因此b的合法范围是：1 ≤ b ≤ min(len(组B), len(组C), 181, 485)

代码实现（Pandas） #

import pandas as pd
import numpy as np

# 拆分原数据集为四个分组
group_a = df[(df['cond_1'] == True) & (df['cond_2'] == True)]
group_b = df[(df['cond_1'] == True) & (df['cond_2'] == False)]
group_c = df[(df['cond_1'] == False) & (df['cond_2'] == True)]
group_d = df[(df['cond_1'] == False) & (df['cond_2'] == False)]

# 验证原数据是否能满足约束
min_b = 1
max_b = min(len(group_b), len(group_c), 181, 485)
if max_b < min_b:
    raise ValueError("原数据中组B或组C样本量不足，无法满足交叉组合要求")

# 随机选择合法的b值（也可指定固定值）
b = np.random.randint(min_b, max_b + 1)
a = 181 - b
c = b
d = 485 - b

# 从各分组抽取对应数量的样本
sample_a = group_a.sample(n=a, random_state=42)
sample_b = group_b.sample(n=b, random_state=42)
sample_c = group_c.sample(n=c, random_state=42)
sample_d = group_d.sample(n=d, random_state=42)

# 合并样本并打乱顺序
df_extracted = pd.concat([sample_a, sample_b, sample_c, sample_d])\
                .sample(frac=1, random_state=42)\
                .reset_index(drop=True)

# 验证结果是否符合约束
print(f"cond_1=True数量: {df_extracted['cond_1'].sum()}")
print(f"cond_2=True数量: {df_extracted['cond_2'].sum()}")
print(f"cond_1=False数量: {len(df_extracted) - df_extracted['cond_1'].sum()}")
print(f"cond_2=False数量: {len(df_extracted) - df_extracted['cond_2'].sum()}")
print(f"总样本数: {len(df_extracted)}")
print(f"包含交叉组合: {(len(sample_b) > 0) or (len(sample_c) > 0)}")

关键注意事项 #

• 若原数据中某分组的样本量小于需要抽取的数量，会抛出错误，需提前确认原数据各分组的规模
• random_state参数用于固定随机种子，方便结果复现，不需要可移除
• 若需多次生成不同子集，可循环执行逻辑并更换随机种子或不设置该参数

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