如何不使用Lambda函数在Python中按分层ID创建聚类抽样样本？

针对百万级数据集的分层有放回聚类抽样（抽n-1个聚类，n为分层内聚类数），原groupby.apply(lambda)方案效率低下，可通过去重分层聚类+向量化抽样大幅提速，具体实现如下：

核心优化思路 #

• 先提取每个分层的唯一聚类（store_id），避免重复处理同一聚类下的多条观测
• 用numpy向量化随机抽样替代逐组lambda调用，消除groupby.apply的性能开销
• 最后将抽样结果与原始数据关联，获取完整观测

优化后代码 #

import pandas as pd
import numpy as np

# 原始数据（示例）
example_data=[['1','A','1',33,23,3,2],['1','A','2',37,20,3,2],['1','A','3',30,27,3,2],['1','A','4',36,21,3,2],['1','A','5',33,23,3,2],['1','B','1',38,20,3,2],\
                     ['1','B','2',39,20,3,2],['1','B','3',33,20,3,2],['1','B','4',33,23,3,2],['1','C','1',27,25,3,2],['1','C','2',28,26,3,2],['2','E','1',38,21,2,1],\
                     ['2','E','2',39,22,2,1],['2','F','1',37,21,2,1],['2','F','2',40,21,2,1],['3','G','1',32,26,4,3],['3','G','2',32,27,4,3],['3','H','1',38,28,4,3],\
                     ['3','H','2',41,28,4,3],['3','H','3',46,22,4,3],['3','H','4',44,23,4,3],['3','H','5',44,28,4,3],['3','H','6',45,30,4,3],['3','I','1',34,29,4,3],\
                     ['3','I','2',32,24,4,3],['3','J','1',25,23,4,3],['3','J','2',21,26,4,3],['3','J','3',22,27,4,3],['4','K','1',20,21,1,1],['4','K','2',24,27,1,1],\
                      ['4','K','3',20,20,1,1]]
df_ex=pd.DataFrame(example_data,columns=['strata','store_id','product','weight','size','stores_in_strata','number_stores_to_sample_from_strata'])

# 步骤1：提取每个分层的唯一store_id及抽样数量
unique_stores = df_ex[['strata', 'store_id', 'number_stores_to_sample_from_strata']].drop_duplicates()

# 步骤2：按分层分组，批量生成抽样结果
sampled_stores = []
for strata, group in unique_stores.groupby('strata'):
    sample_size = group['number_stores_to_sample_from_strata'].iloc[0]
    # 用numpy随机抽样（有放回），比pandas.sample更高效
    sampled_idx = np.random.choice(group.index, size=sample_size, replace=True)
    sampled_stores.append(group.loc[sampled_idx])

# 合并抽样结果
sampled_df = pd.concat(sampled_stores, ignore_index=True)

# 步骤3：关联原始数据，获取完整观测
rth_samp = pd.merge(df_ex, sampled_df[['strata', 'store_id']], on=['strata', 'store_id'], how='inner')

性能提升关键点 #

• 减少重复计算：仅对每个分层的唯一store_id抽样，避免处理同一store下的多条product记录
• 向量化抽样：用numpy.random.choice替代groupby.apply(lambda)，消除逐组循环的Python层开销
• 简化合并逻辑：仅关联必要的strata和store_id字段，减少合并数据量

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