如何使用Python将多变量数据集拆分为分布一致的训练集和测试集

你提出的拆分方案完全可行，这就是典型的分层随机拆分逻辑，刚好适配你要保证各亚型在数据集占比一致的需求。

由于你当前数据集不同原发肿瘤位点的样本量差异较大：比如乳腺来源有5份样本，皮肤、骨髓来源仅各有2份样本，如果直接对全量数据集做随机拆分，很容易出现小样本亚型的样本全部流入训练集或测试集的情况，导致两个数据集的分布不一致。
你先按原发位点亚型拆分、每个亚型单独做50/50随机划分，再分别合并训练、测试子集的操作，刚好可以规避上述问题，能保证最终得到的训练集、测试集里各个原发位点的样本占比和原始数据集完全一致，同时两者的总样本量也会基本接近。

如果碰到样本量为奇数、无法完全均分的亚型（比如你示例中的乳腺来源共5份样本），可以固定随机种子后，轮流将多出来的1份样本划分到训练集/测试集即可，不会对整体分布一致性产生明显影响。

如果想要简化操作，也可以直接调用成熟的机器学习工具库的分层拆分接口，不需要手动逐个亚型拆分。以Python的scikit-learn库为例，一行代码即可实现等价逻辑：

from sklearn.model_selection import train_test_split
# X为所有影像路径列表，y为每个影像对应的原发位点标签列表
train_X, test_X, train_y, test_y = train_test_split(X, y, test_size=0.5, stratify=y, random_state=42)

其中stratify=y参数会自动按照标签的类别分布做分层拆分，和你手动拆分的效果完全一致。

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