Python实现带先验比例的分层抽样函数

分层抽样的常见应用场景是选取能大致保留所选变量分布的随机样本，以此确保样本具备代表性。

需要创建一个分层抽样函数，不使用原数据集的比例，而是按照指定的目标变量比例来完成抽样。

def stratified_sampling_prior(df, column, prior_dict, sample_size):
    ...
    return df_sampled

参数说明 #

• df：输入的数据集
• column：用于执行分层抽样的分类变量
• prior_dict：字典类型，包含目标变量各分类的抽样占比（以小数形式表示）
• sample_size：期望得到的样本总实例数量

以下是可运行的示例代码：

import pandas as pd

prior_dict = {
    "A": 0.2,
    "B": 0.2,
    "C": 0.1,
    "D": 0.5
}

df = pd.DataFrame({
    "Category": ["A"]*10 + ["B"]*50 + ["C"]*15 + ["D"]*100,
    "foo": ["foo" for i in range(175)],
    "bar": ["bar" for i in range(175)]
})

若采用传统分层抽样（遵循原数据集比例），各分类占比如下：

df["Category"].value_counts()/df.shape[0]*100
# 输出结果：
# D    57.14
# B    28.57
# C     8.57
# A     5.71

而使用prior_dict指定的比例抽样后，样本的分类占比预期为：

df_sample = stratified_sampling_prior(df, "Category", prior_dict, sample_size=100)
df_sample["Category"].value_counts()/df_sample.shape[0]*100
# 输出结果：
# D    50.00
# B    20.00
# C    10.00
# A    20.00

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