python - 根据列的值将数据框拆分为两个文件

Question

我需要将数据框分成两部分。例如，如果以下数据帧基于 Col1 随机拆分，则两个文件都应包含来自每个类别 1,2 和 3 的样本。

Col1 col2 
1      a
1      b
2      c
2      d
3      e

到目前为止，我可以使用sklearn.cross_validation import train_test_split. 但是我无法弄清楚应该如何进行拆分以从每个类别中提取样本。 所有帮助将不胜感激。谢谢。

Answer 1

看看StratifiedKFold对象。

http://scikit-learn.org/stable/modules/generated/sklearn.cross_validation.StratifiedKFold.html

文档中有一个简短的示例，展示了如何使用它