如何从Pandas DataFrame抽取满足多变量条件的随机样本？

首先还原你的数据集：

import pandas as pd

data = {'character': {0: 'mario',
  1: 'luigi',
  2: 'yoshi',
  3: 'peach',
  4: 'bowser',
  5: 'boo',
  6: 'toad',
  7: 'lakitu',
  8: 'koopa'},
 'location': {0: 'castle',
  1: 'castle',
  2: 'forest',
  3: 'castle',
  4: 'dungeon',
  5: 'dungeon',
  6: 'forest',
  7: 'forest',
  8: 'dungeon'},
 'color': {0: 'red',
  1: 'green',
  2: 'green',
  3: 'yellow',
  4: 'blue',
  5: 'white',
  6: 'blue',
  7: 'yellow',
  8: 'green'}}

df1 = pd.DataFrame(data)

关于单一函数实现的说明 #

Pandas内置的df.sample()函数无法直接满足你提出的双重约束抽样需求（每个location至少2个角色、每个color至少1个角色），因为它仅支持简单随机抽样或基于权重的抽样，没有内置分层多条件约束的逻辑。

不过你可以把满足约束的抽样逻辑封装成一个自定义函数，实现类似“单一函数调用”的效果，具体实现步骤如下：

自定义约束抽样函数实现 #

def constrained_sample(df, location_min=2):
    # 第一步：按location分组，每组抽取至少location_min个样本
    loc_samples = df.groupby('location').sample(n=location_min, random_state=42)
    
    # 第二步：检查已抽取样本是否覆盖所有color类别
    existing_colors = loc_samples['color'].unique()
    missing_colors = df[~df['color'].isin(existing_colors)]['color'].unique()
    
    # 第三步：补充缺失的color样本（每个缺失color抽1个）
    if len(missing_colors) > 0:
        color_supplements = df[df['color'].isin(missing_colors)].groupby('color').sample(n=1, random_state=42)
        final_sample = pd.concat([loc_samples, color_supplements]).drop_duplicates()
    else:
        final_sample = loc_samples
    
    return final_sample

# 调用函数获取满足条件的样本
result = constrained_sample(df1)

逻辑解释 #

1. 按location分层抽样：先确保每个location分组都能拿到至少2个随机样本，满足第一个约束条件。
2. 检查color覆盖：判断当前样本是否包含所有color类别，如果有缺失则从原数据中抽取对应color的样本补充。
3. 去重处理：避免补充样本和已抽样本重复（比如某个样本同时属于需要补充的color和已抽的location组）。

示例输出说明 #

以random_state=42为例，第一次按location抽取的样本可能是：

这个样本已经覆盖了所有color（red、yellow、green、blue、white），因此不需要补充，直接作为最终样本。如果某次抽样未覆盖全部color，函数会自动补充对应样本。

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