Pandas + Scikit 学习 : issue with stratified k-fold

Question

当与数据框一起使用时，StratifiedKFold from scikit-learn 返回从 0 到 n 的索引列表，而不是来自 DF 索引的值列表。有没有办法改变它？

前任 :

df = pd.DataFrame()
df["test"] = (0, 1, 2, 3, 4, 5, 6)
df.index   = ('a', 'b', 'c', 'd', 'e', 'f', 'g')
for i, (train, test) in enumerate(StratifiedKFold(df.index)):
    print i, (train, test)

给出:

0 (array([], dtype=64), array([0,1,2,3,4,5,6])
1 (array([0,1,2,3,4,5,6]), array([], dtype=64))
2 (array([0,1,2,3,4,5,6]), array([], dtype=64))

我希望返回 df 的索引，而不是 df 的长度范围...

Answer 1

你得到的数字只是 df.index 的索引选择者 StratifiedKFold .

要将其改回 DataFrame 的索引，只需

for i, (train, test) in enumerate(StratifiedKFold(df.index)):
    print i, (df.index[train], df.index[test])

这使

0 (Index([], dtype='object'), Index([u'a', u'b', u'c', u'd', u'e', u'f', u'g'], dtype='object'))
1 (Index([u'a', u'b', u'c', u'd', u'e', u'f', u'g'], dtype='object'), Index([], dtype='object'))
2 (Index([u'a', u'b', u'c', u'd', u'e', u'f', u'g'], dtype='object'), Index([], dtype='object'))