实现的逻辑比较简单,只是把dataframe自带的sample函数结合了value_counts()进行了优化,从而实现了分层抽样。
import pandas
#df_data:输入的数据
#stratify:分层参照的标签列名
#proportion:抽样的比例数,默认30%
def stratified_sampling(df_data, stratify, proportion = 0.3):
vc = df_data[stratify].value_counts()
sam = pd.DataFrame(columns = df_data.columns.tolist())
for vi in vc.index:
dd = df_data[df_data[stratify] == vi].sample(n = round(vc[vi] * proportion))
sam = pd.concat([sam, dd ], ignore_index = False)
return sam