根据列的值将数据帧拆分为两个文件

时间:2016-06-07 11:34:24

标签: python sklearn-pandas

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

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

到目前为止,我可以使用sklearn.cross_validation import train_test_split将数据拆分为所需的比例。 但我无法弄清楚应如何分割来从每个类别中提取样本。 所有帮助将不胜感激。感谢。

1 个答案:

答案 0 :(得分:1)

看看StratifiedKFold对象。

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

文档中有一个简短示例,说明如何使用它