Question

我有一个数据框，让我们说：

df = pd.DataFrame({'id': [1, 1 , 2, 2, 2, 3, 4, 5], 'val1': [ 1, 2, 1, 1, 2, 1, 2, 3], 'val2': [3, 3, 4, 4, 4, 3, 4, 4]})

我想使用train列中的值将其拆分为两个数据框（test和id）。拆分应该是这样的，在第一个数据帧中我有80％的（唯一的）id s，而在第二个数据帧中，我有20％的id s。 id应随机分割。

我自己的尝试：

import random
import pandas as pd
def train_test_split(df, test_size=0.2, prng_seed=None):
    prng = random.Random()
    prng.seed(prng_seed)
    id_list = df['id'].unique().tolist()
    prng.shuffle(id_list)
    id_size = len(id_list)
    test_abs_size =  int(id_size * test_size)
    test_id = id_list[-test_abs_size:]
    train_id = id_list[:-test_abs_size]
    train_data = df[df['id'].isin(train_id)]
    test_data = df[df['id'].isin(test_id)]
    return train_data, test_data

Answer 1

以下代码将数据集拆分为80-20列车测试集

import pandas as pd
import numpy as np
from sklearn.cross_validation import train_test_split

train, test = train_test_split(df, test_size = 0.2)

根据@JanTrienes评论，如果您想保留id的分发，可以使用stratify以下代码表示：

import pandas as pd
from sklearn.cross_validation import train_test_split

df = pd.DataFrame({'id': [1, 1, 2, 2, 2, 3, 4, 4,
                          1, 1, 2, 2, 2, 3, 4, 4],
                 'val1': [1, 2, 1, 1, 2, 1, 2, 3,
                          1, 2, 1, 1, 2, 1, 2, 3],
                 'val2': [3, 3, 4, 4, 4, 3, 4, 4,
                          3, 3, 4, 4, 4, 3, 4, 4]})

train, test = train_test_split(df, test_size = 0.2, stratify=df.id)

以下是输出结果的示例：

train:
    id  val1  val2
0    1     1     3
7    4     3     4
15   4     3     4
13   3     1     3
14   4     2     4
11   2     1     4
9    1     2     3
8    1     1     3
12   2     2     4
4    2     2     4
2    2     1     4
5    3     1     3
test:
    id  val1  val2
6    4     2     4
10   2     1     4
1    1     2     3
3    2     1     4

根据列中的值将pandas数据帧拆分为两个数据帧

1 个答案: