Python有放回采样时限制每类别采样数量的可用库函数查询

场景1：分层采样需求 #

如果你的使用场景是分层采样，需要按子群体限制采样次数，以下两个常用库的现成函数可以直接使用：

• scikit-learn的sklearn.model_selection.train_test_split方法支持stratify参数，传入分层标签后会自动按各子群体的占比做无放回采样，天然避免采样次数超过子群体实际成员数的问题，适合数据集拆分场景。
• pandas的DataFrame.sample方法原生支持weights参数指定采样权重，设置replace=False即为无放回采样，保证每个样本最多被抽中一次。如果需要按子群体设置独立的采样数量，可先按子群体分组再对每个分组单独采样，示例代码如下：

import pandas as pd

# 示例数据：包含子群体标签group、样本权重weight
df = pd.DataFrame({
    "value": range(100),
    "group": ["A"]*30 + ["B"]*50 + ["C"]*20,
    "weight": [i/100 for i in range(100)]
})

# 按group分组，每个组最多采样10个，带权重无放回
sampled_df = df.groupby("group", group_keys=False).apply(
    lambda x: x.sample(n=min(10, len(x)), replace=False, weights="weight")
)

场景2：通用带权重、自定义单元素最大采样次数 #

如果需要自定义每个元素的最大允许出现次数（不限于无放回的1次限制），有两种成熟实现方式：

标准库实现（Python 3.9+） #

Python 3.9及以上版本的标准库random模块已经原生支持该能力，random.sample方法新增了counts参数和weights参数，可直接指定每个元素的最大出现次数和采样权重，不需要额外引入第三方库，示例代码如下：

import random

# 示例：采样总体为[a,b,c]，对应权重1、2、3，对应最大允许出现次数2、1、3，总共采样5个
population = ["a", "b", "c"]
sample_weights = [1, 2, 3]
max_occur_counts = [2, 1, 3]
result = random.sample(population, k=5, counts=max_occur_counts, weights=sample_weights)

Numpy高性能实现 #

如果需要处理大规模数据，可通过Numpy扩展数组的方式实现，不需要手动归一化权重也能快速完成采样：

import numpy as np

population = np.array(["a", "b", "c"])
weights = np.array([1,2,3])
max_counts = np.array([2,1,3])

# 按最大出现次数扩展总体和对应权重
extended_pop = np.repeat(population, max_counts)
extended_weights = np.repeat(weights, max_counts)
# 自动归一化权重为概率分布
p = extended_weights / extended_weights.sum()
# 无放回采样5个
result = np.random.choice(extended_pop, size=5, replace=False, p=p)

如果使用的Python版本低于3.9且不想依赖第三方库，也可以参考上述Numpy的思路，手动扩展总体列表后调用旧版本的random.sample完成采样。

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