Python pandas 分层抽样 超简洁

这段时间在打一个机器学习比赛，遇到分层抽样的需求。在网上查了一下，虽然也有前辈用pandas写出来过，但代码有很多冗余，于是我又重写了一个超简洁的版本。注释应该还算清晰，各位可以自取。用得好可以给个赞，感谢！

sample_data：原始数据，示例中‘label’列是分层依据
train_proportion：训练集抽取比例，默认70%
seed：随机抽取与打乱顺序时用到的随机数种子，默认为None。当输入一个数字的时候，输出的样本及样本顺序都会固定。

返回的是训练样本和测试样本，两个样本不重叠，且合在一起包含了全部样本。

def stratified_sampling(sample_data, train_proportion=0.7, seed=None):
    # 分层随机选取训练样本
    train_sample = sample_data.groupby('label').apply(lambda x: x.sample(frac=train_proportion, random_state=seed))
    # 上一步得到的样本是按照label聚在一起的，所以这一步需要打乱顺序
    train_sample = train_sample.sample(frac=1, random_state=seed)
    train_sample.reset_index(drop=True, inplace=True)
    # 和原始数据拼接起来，删掉重复值，从而获得测试数据
    test_sample = pd.concat([sample_data, train_sample]).drop_duplicates(keep=False)
    # 同样对抽取出的测试样本打乱顺序
    test_sample = test_sample.sample(frac=1, random_state=seed)
    test_sample.reset_index(drop=True, inplace=True)
    return train_sample, test_sample

然而作者写完这段代码才发现sklearn.model_selection里面有train_test_split，可以完全实现上述效果。不过还是需要reset_index

data_train, data_test = train_test_split(data_label, test_size=0.3, random_state=0, stratify=data_label.label)