函数实现将 DataFrame 数据直接划分为测试集训练集

 虽然 Scikit-Learn 有可以划分数据集的函数 train_test_split ，但在有些特殊情况我们只希望它将 DataFrame 数据直接划分为 train, test 而不是像 train_test_split 返回四个值。这里写了一个类似功能的函数：

import numpy as np
import pandas as pd
from sklearn.utils import shuffle as resetdef train_test_split(data, test_size=0.3, shuffle=True, random_state=None):
'''Split DataFrame into random train and test subsets Parameters
----------
data : pandas dataframe, need to split dataset. test_size : float
If float, should be between 0.0 and 1.0 and represent the
proportion of the dataset to include in the train split. random_state : int, RandomState instance or None, optional (default=None)
If int, random_state is the seed used by the random number generator;
If RandomState instance, random_state is the random number generator;
If None, the random number generator is the RandomState instance used
by `np.random`. shuffle : boolean, optional (default=None)
Whether or not to shuffle the data before splitting. If shuffle=False
then stratify must be None.
''' if shuffle:
data = reset(data, random_state=random_state) train = data[int(len(data)*test_size):].reset_index(drop = True)
test = data[:int(len(data)*test_size)].reset_index(drop = True) return train, test


效果如下：