基于Pandas的多层加权Bootstrap抽样问题

需求说明 #

• 对包含Country、Fruit（及Vegetable）列的数据集执行Bootstrap抽样
• 核心要求：Country与Fruit/Vegetable独立均匀随机抽样
• 额外约束：同一Bootstrap迭代中，Fruit和Vegetable的抽样Country集合必须完全一致，且避免使用嵌套循环

递进抽样策略 #

针对需求设计了4种逐步接近目标的抽样方案：

1. 每个Country均匀抽取M个Fruit #

对每个Country独立均匀抽取指定数量的Fruit，不涉及Country层面的Bootstrap抽样。

sampled_fruits = df.groupby('Country').sample(n=M, replace=True, random_state=rng)

2. 每个Country按Fruit出现次数反比抽取M个Fruit #

在每个Country内，按Fruit全局出现次数的反比权重抽取样本，解决稀有Fruit抽样不足的问题。

df2 = df.join(fruit_weights, on='Fruit')  # 为原数据添加权重列
sampled_fruits = df2.groupby('Country').sample(
    n=M,
    replace=True,
    random_state=rng,
    weights='fruit_weights',
)

3. 整体均匀抽取（平均）M个Fruit，对Country进行Bootstrap抽样 #

先对Country进行Bootstrap抽样，再基于抽样结果抽取Fruit，实现Country层面的Bootstrap，但Fruit抽样为均匀分布。

sampled_fruits = pd.concat(
    {
        s: df.sample(
            n=df.index.nunique(),  # 抽样数量等于Country的唯一数量
            weights=country_weights,
            replace=True,
            random_state=rng,
        )
        for s in range(M)
    },
    names=['sample', 'Country'],
).reset_index('sample')

4. 按Fruit出现次数反比抽取（平均）M个Fruit，同时对Country进行Bootstrap抽样 #

结合Country的Bootstrap抽样和Fruit的反比权重抽样，接近目标但存在多类别抽样时Country集合不一致的问题。

df4 = df.join(fruit_weights, on='Fruit')

# 按Country归一化Fruit权重，避免影响Country权重的作用
df4['fruit_weights'] = df4.fruit_weights.groupby('Country').transform(lambda x: x / x.sum())

df4 = df4.join(country_weights, on='Country')

# 计算联合权重
weight_cols = [c for c in df4.columns if '_weights' in c]
weights = df4[weight_cols].prod(axis=1)
df4 = df4.drop(columns=weight_cols)

sampled_fruits = pd.concat(
    {
        s: df4.sample(
            n=df.index.nunique(),
            weights=weights,
            replace=True,
            random_state=rng,
        )
        for s in range(M)
    },
    names=['sample', 'Country'],
).reset_index('sample')

现有方案的问题 #

策略4已接近需求，但当同时对Vegetable（其类别数量、按Country划分的分布均与Fruit不同）进行抽样时，同一Bootstrap迭代中，Fruit和Vegetable的抽样Country集合不一致——因为每个类别单独执行了Country的Bootstrap抽样，导致迭代间的Country集合不匹配。

优化方案：统一Country集合的多类别Bootstrap抽样 #

核心思路 #

先预先生成所有Bootstrap迭代的Country抽样集合，再基于该集合分别对Fruit和Vegetable进行抽样，确保同一迭代的Country集合完全一致，同时避免嵌套循环。

代码实现 #

扩展示例数据（包含Vegetable列）

import pandas as pd
import numpy as np

# 原Fruit数据集
df_fruit = pd.DataFrame(
    np.array([
        ['USA', 'USA', 'Canada', 'Canada', 'Canada', 'Canada', 'Canada', 'Mexico', 'Mexico', 'Mexico', 'Mexico', 'Brazil', 'Brazil', 'Brazil', 'Brazil', 'UK', 'UK', 'France', 'France', 'Germany', 'Germany', 'Germany', 'Germany', 'Germany', 'Italy', 'Italy', 'Spain', 'Spain', 'Spain', 'Spain', 'Spain'],
        ['Pineapple', 'Apple', 'Pineapple', 'Apple', 'Cherry', 'Watermelon', 'Orange', 'Apple', 'Banana', 'Cherry', 'Orange', 'Watermelon', 'Banana', 'Apple', 'Blueberry', 'Cherry', 'Apple', 'Banana', 'Blueberry', 'Banana', 'Apple', 'Cherry', 'Blueberry', 'Pineapple', 'Pineapple', 'Watermelon', 'Pineapple', 'Watermelon', 'Apple', 'Orange', 'Blueberry'],
    ]).T,
        columns=['Country', 'Fruit'],
).set_index('Country')
df_fruit['other_fruit_cols'] = '...'

# 新增Vegetable数据集（模拟与Fruit不同的分布）
df_veg = pd.DataFrame(
    np.array([
        ['USA', 'Canada', 'Canada', 'Mexico', 'Mexico', 'Mexico', 'Brazil', 'Brazil', 'UK', 'UK', 'UK', 'France', 'Germany', 'Germany', 'Germany', 'Italy', 'Spain', 'Spain', 'Spain'],
        ['Carrot', 'Potato', 'Broccoli', 'Carrot', 'Potato', 'Broccoli', 'Carrot', 'Potato', 'Carrot', 'Potato', 'Broccoli', 'Potato', 'Carrot', 'Potato', 'Broccoli', 'Carrot', 'Potato', 'Broccoli', 'Carrot'],
    ]).T,
        columns=['Country', 'Vegetable'],
).set_index('Country')
df_veg['other_veg_cols'] = '...'

配置参数

M = 10  # Bootstrap迭代次数
NUM_COUNTRIES_PER_SAMPLE = df_fruit.index.nunique()  # 每次抽样的Country数量
rng = np.random.default_rng(seed=123)  # 随机种子确保可复现

# 生成Fruit和Vegetable的权重（反比于出现次数）
fruit_weights = 1 / df_fruit.groupby('Fruit').size().rename('fruit_weights')
fruit_weights /= fruit_weights.sum()

veg_weights = 1 / df_veg.groupby('Vegetable').size().rename('veg_weights')
veg_weights /= veg_weights.sum()

# Country权重（反比于出现次数）
country_weights = 1 / pd.concat([df_fruit.index, df_veg.index]).value_counts().rename('country_weights')
country_weights /= country_weights.sum()

预生成统一的Country抽样集合

# 生成M次迭代的Country抽样结果，形状为(M, NUM_COUNTRIES_PER_SAMPLE)
sampled_countries = rng.choice(
    country_weights.index,
    size=(M, NUM_COUNTRIES_PER_SAMPLE),
    p=country_weights.values,
    replace=True
)

# 转换为DataFrame，标记迭代编号（sample_id）和Country
country_samples_df = pd.DataFrame({
    'sample_id': np.repeat(range(M), NUM_COUNTRIES_PER_SAMPLE),
    'Country': sampled_countries.flatten()
}).set_index(['sample_id', 'Country'])

基于统一Country集合抽取Fruit和Vegetable

# 处理Fruit数据集：添加权重，合并Country抽样集合，按迭代和Country抽样
df_fruit_weighted = df_fruit.join(fruit_weights, on='Fruit')
# 按Country归一化Fruit权重
df_fruit_weighted['fruit_weights'] = df_fruit_weighted['fruit_weights'].groupby('Country').transform(lambda x: x/x.sum())

# 合并预生成的Country抽样集合，确保每个迭代的Country都对应到Fruit数据
fruit_samples = country_samples_df.join(df_fruit_weighted, on='Country')
# 按sample_id和Country分组，每组抽取1个Fruit（可根据需求调整数量）
fruit_sampled_final = fruit_samples.groupby(['sample_id', 'Country']).sample(
    n=1,
    weights='fruit_weights',
    replace=True,
    random_state=rng
).reset_index('Country')

# 处理Vegetable数据集，逻辑完全一致
df_veg_weighted = df_veg.join(veg_weights, on='Vegetable')
df_veg_weighted['veg_weights'] = df_veg_weighted['veg_weights'].groupby('Country').transform(lambda x: x/x.sum())

veg_samples = country_samples_df.join(df_veg_weighted, on='Country')
veg_sampled_final = veg_samples.groupby(['sample_id', 'Country']).sample(
    n=1,
    weights='veg_weights',
    replace=True,
    random_state=rng
).reset_index('Country')

验证一致性

# 检查同一迭代的Country集合是否一致
assert all(fruit_sampled_final.groupby('sample_id').apply(lambda x: set(x.index)) == veg_sampled_final.groupby('sample_id').apply(lambda x: set(x.index)))

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