Python 自动将数据按比例随机为训练集和测试集代码

PART 1. 想在网上找到关于将数据集自动抽样的方法，看见了其实Pandas库下有一个方法可以抽样，但是感觉不是很灵活，而且数据类型是dataframe的，不是很习惯，这里就写了一个方法，输入是数据集X（m*n），数据集类别标签Y（m*1），训练集比例pct（0～1），返回训练数据集和测试数据集，返回的数据中最后一列是类别标签。

import random
import numpy as np
def seperateData(x,y,pct):
    # x：数据
    # y：类别标签
    # pct：训练集所占比例：区间为0~1
    data = np.column_stack((x,y))

    length = len(y)
    # index = round(length*pct) # (之前是这一行，删掉)
    index = int(round(length*pct))
    
    indexes = np.array(range(0,length))
    random.shuffle(indexes)
    trn_idxes = indexes[0:index]
    tst_idxes = indexes[index:length]
    
    trn = data[trn_idxes,:]
    tst = data[tst_idxes,:]
    return trn, tst


# 测试样例：
x = pd.DataFrame(np.random.rand(10,10))
y = pd.DataFrame(np.sign(np.random.rand(10)-0.5))
seperateData(x,y,0.3)

 

PART 2.事实是，忽然发现确实有个函数可以自动生成训练集和测试集的，但是没有用pandas下的shuffle，只使用了简单的sklearn的函数

from numpy import random
import numpy as np

# （1）可选：数据为数组形式
x = np.floor(10*random.rand(num,2))
y = np.floor(10*random.rand(num,1))

# （2）可选：数据为dataframe形式
x = pd.DataFrame(np.random.rand(10,10))
y = pd.DataFrame(np.sign(np.random.rand(10)-0.5))

from sklearn import cross_validation

X_train, X_test, y_train, y_test = cross_validation.train_test_split(x, y, test_size=0.4, random_state=0)

 

好吧，get了两个方法！