在机器学习中,按照分类列拆分训练集和测试集可以使用以下方法:
from sklearn.model_selection import train_test_split
# 假设data是包含特征和目标变量的DataFrame,其中category是分类列
X_train, X_test, y_train, y_test = train_test_split(data.drop('target', axis=1), data['target'], test_size=0.2, stratify=data['category'])
import pandas as pd
# 假设data是包含特征和目标变量的DataFrame,其中category是分类列
train_data = data.groupby('category').apply(lambda x: x.sample(frac=0.8, random_state=42))
test_data = data.drop(train_data.index)
import numpy as np
# 假设data是包含特征和目标变量的numpy数组,其中category是分类列
unique_categories = np.unique(data[:, -1]) # 获取唯一的分类值
train_data, test_data = [], []
for category in unique_categories:
category_data = data[data[:, -1] == category]
np.random.shuffle(category_data) # 随机打乱数据
train_size = int(len(category_data) * 0.8)
train_data.append(category_data[:train_size])
test_data.append(category_data[train_size:])
train_data = np.vstack(train_data)
test_data = np.vstack(test_data)
以上代码示例分别使用了sklearn库的train_test_split函数、pandas库的groupby和apply方法,以及numpy库的unique函数和数组切片操作,根据分类列将数据集拆分为训练集和测试集。具体选择哪种方法取决于你使用的库和数据结构。