Pandas数据帧分层分解为训练集、验证集和测试集

Question

以下极度简化的数据框表示包含医疗诊断的更大数据框：

medicalData = pd.DataFrame({'diagnosis':['positive','positive','negative','negative','positive','negative','negative','negative','negative','negative']})
medicalData

    diagnosis
0   positive
1   positive
2   negative
3   negative
4   positive
5   negative
6   negative
7   negative
8   negative
9   negative

问题：对于机器学习，我需要按以下方式将此数据帧随机拆分为三个子帧：

trainingDF, validationDF, testDF = SplitData(medicalData,fractions = [0.6,0.2,0.2])

…其中分割数组指定进入每个子帧的完整数据的分数

• 子帧中的数据需要是互斥的，分割数组（分数）需要求和为一
• 传统上，每个子集的阳性诊断分数需要大致相同。
• Answers to this question 建议使用the pandas sample method或the train_test_split function from sklearn。但这些解决方案似乎都不能很好地推广到n个分裂，也没有一个提供分层分裂

Answer 1

纯pandas溶液

按70/20/10%的比例分为培训/验证/测试：

train_df = df.sample(frac=0.7, random_state=random_seed)
tmp_df = df.drop(train_df.index)
test_df = tmp_df.sample(frac=0.33333, random_state=random_seed)
valid_df = tmp_df.drop(test_df.index)

assert len(df) == len(train_df) + len(valid_df) + len(test_df), "Dataset sizes don't add up"
del tmp_df