[自用] label添加到feature模板

from sklearn.model_selection import StratifiedShuffleSplit
import numpy as np
X = np.array([[111] , [112] , [113] ,[114],  [222] , [221] , [223] ,[224], [331] , [332] , [333] ,[334] , [441] , [442] , [443], [444]])
y = np.array([1, 1, 1,1 , 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4]) #label
y = y.reshape(-1,1)#要给它转换成列向量
data_withlabel = np.hstack((X, y))#这是把y加到X的右边
train_indices = []
test_indices = []

split = StratifiedShuffleSplit(n_splits=1, test_size=0.25, random_state=0)


for train_index, test_index in split.split(data_withlabel, data_withlabel[:,-1]):
    train_indices.extend(train_index)
    test_indices.extend(test_index)
print(test_index.size)
# 根据索引创建训练和测试数据集
X_train = X[train_indices]
X_test = X[test_indices]
y_train = y[train_indices]
y_test = y[test_indices]
train_dataset = data_withlabel[train_indices]
test_dataset = data_withlabel[test_indices]

print(train_dataset[:, -1])
print(test_dataset)