Pandas Concat未正确合并DataFrame：鸢尾花训练集维度异常问题

问题根源 #

鸢尾花数据集每个类别只有50个样本，你设置test_size=0.8意味着只保留20%的数据作为训练集，每个类别训练集数量是50*0.2=10，三个类别加起来就是30，所以合并后得到30×4的结果，完全不符合你要每个类别40个训练样本的需求。

VSCode调试的维度矛盾大概率是因为索引重复或者调试面板的缓存显示问题，修正代码后就能解决。

修正方案 #

把train_test_split的test_size改成0.2，这样每个类别会保留80%（40个）样本作为训练集，三个类别合并后就是120×4的预期维度。同时可以简化代码，不用单独提取索引再筛选，直接按类别拆分更高效：

from sklearn import datasets
from sklearn.model_selection import train_test_split
import pandas as pd

# 加载数据集
iris = datasets.load_iris()
x = pd.DataFrame(data=iris['data'], columns=iris['feature_names'])
y = pd.DataFrame(data=iris['target'], columns=['target'])

# 按类别拆分数据
iris_0 = x[y['target'] == 0]
iris_1 = x[y['target'] == 1]
iris_2 = x[y['target'] == 2]

# 拆分训练集（每个类别取40个，test_size=0.2即测试集10个，训练集40个）
x_train_0, x_test_0, _, _ = train_test_split(iris_0, y[y['target'] == 0], test_size=0.2, random_state=42)
x_train_1, x_test_1, _, _ = train_test_split(iris_1, y[y['target'] == 1], test_size=0.2, random_state=42)
x_train_2, x_test_2, _, _ = train_test_split(iris_2, y[y['target'] == 2], test_size=0.2, random_state=42)

# 合并训练集
x_train = pd.concat([x_train_0, x_train_1, x_train_2])
# 验证维度
print(x_train.shape)  # 输出 (120, 4)

额外说明 #

• 添加random_state=42可以保证每次拆分结果一致，方便调试和复现
• 不需要的变量可以用_接收，避免冗余

内容的提问来源于stack exchange，提问作者Watson221