
本文介绍一种灵活、可复用的方法,通过分组抽样与去重合并策略,在 pandas 中实现同时满足“每类位置至少2个样本”和“每类颜色至少1个样本”的多维约束随机抽样。
本文介绍一种灵活、可复用的方法,通过分组抽样与去重合并策略,在 pandas 中实现同时满足“每类位置至少2个样本”和“每类颜色至少1个样本”的多维约束随机抽样。
在实际数据分析中,简单的 df.sample(n) 往往无法满足业务逻辑所需的结构化抽样要求。例如,当需要确保样本在多个分类维度(如 location 和 color)上均具备最低覆盖度时,就必须采用分层+组合的策略。以下是一个专业、健壮且易于扩展的解决方案:
核心思路:分组抽样 → 合并去重 → 重置索引
我们分别对两个关键列进行约束抽样:
-
按
location抽样:每个位置类别中最多取 2 条记录(若该组不足 2 条,则全取),使用min(2, len(g))保证安全性; -
按
color抽样:每个颜色类别中强制取 1 条(只要该组非空),满足“至少 1 个”的硬性要求; - 最后将两批样本拼接(
pd.concat),并通过.drop_duplicates()去除重复行(避免同一行因同时满足两个条件而被重复纳入),再重置索引以获得整洁结果。
完整可运行代码示例:
import pandas as pd
# 构建原始数据
df1 = pd.DataFrame({
'character': ['mario', 'luigi', 'yoshi', 'peach', 'bowser', 'boo', 'toad', 'lakitu', 'koopa'],
'location': ['castle', 'castle', 'forest', 'castle', 'dungeon', 'dungeon', 'forest', 'forest', 'dungeon'],
'color': ['red', 'green', 'green', 'yellow', 'blue', 'white', 'blue', 'yellow', 'green']
})
def custom_sample(df, n_per_location=2, n_per_color=1, random_state=42):
"""
多条件约束随机抽样函数
Parameters:
-----------
df : pd.DataFrame
输入数据框
n_per_location : int
每个 location 类别最多抽取数量(默认2)
n_per_color : int
每个 color 类别最少抽取数量(默认1)
random_state : int
随机种子,确保结果可复现
Returns:
--------
pd.DataFrame : 满足双重约束的去重合并样本
"""
# Step 1: 按 location 抽样(保底 + 上限控制)
location_samples = (
df.groupby('location', group_keys=False)
.apply(lambda g: g.sample(n=min(n_per_location, len(g)),
random_state=random_state))
)
# Step 2: 按 color 抽样(强制每类至少1条)
color_samples = (
df.groupby('color', group_keys=False)
.apply(lambda g: g.sample(n=min(n_per_color, len(g)),
random_state=random_state))
)
# Step 3: 合并并去重(关键!防止同一行被重复选中)
combined_sample = pd.concat([location_samples, color_samples]).drop_duplicates()
return combined_sample.reset_index(drop=True)
# 调用示例
result = custom_sample(df1)
print(result)
注意事项与优化建议:
- ✅ 可复现性:显式设置
random_state是关键,便于调试与协作; - ✅ 鲁棒性:
min(n, len(g))自动适配稀疏类别(如某 color 只有 1 行,则不会报错); - ⚠️ 样本量不可控:最终样本数 ≈
sum(min(2, size_per_loc)) + sum(min(1, size_per_col)),但经去重后可能略少;如需严格控制总数,可在返回前添加.sample(n=target_size, replace=False)(注意可能破坏约束,需谨慎); - ? 扩展性:如新增约束(如按
character首字母分组),只需增加对应groupby().apply(...)分支并加入concat即可; - ? 验证建议:抽样后推荐执行断言检查,例如:
assert all(result.groupby('location').size() >= 2) or len(result['location'].unique()) == 1 assert all(result.groupby('color').size() >= 1)
该方法平衡了简洁性、可读性与工程实用性,是处理多维分层抽样的典型范式。










