Python中*args解包后，如何重新打包train_test_split的返回结果？

当你用*x把列表里的多个输入传入train_test_split后，返回的结果是按「每个输入的训练集→测试集」的顺序排列的，最后才是目标变量y的训练集和测试集。要重新打包成和原x结构对应的训练/测试集合，可以按以下步骤操作：

核心思路 #

假设原列表x有n个元素（比如矩阵、标量、其他自定义数据），那么train_test_split(*x, y)返回的结果结构是：
[x0_train, x0_test, x1_train, x1_test, ..., xn-1_train, xn-1_test, y_train, y_test]

我们需要把这些结果拆分成三部分：

• 所有x分量的训练集（和原x结构完全对应）
• 所有x分量的测试集（和原x结构完全对应）
• y的训练集与测试集

代码实现 #

import numpy as np
from sklearn.model_selection import train_test_split

# 模拟你的输入结构
x = [np.random.rand(100, 5), np.random.rand(100), np.random.randint(0, 2, 100)]
y = np.random.randint(0, 2, 100)

# 执行拆分
split_results = train_test_split(*x, y, test_size=0.2, random_state=42, shuffle=True)

# 计算原x的元素个数
n = len(x)

# 分离x相关的拆分结果和y相关的拆分结果
x_split_parts = split_results[:2*n]
y_split_parts = split_results[2*n:]

# 重新打包训练集和测试集的x部分
train_x = x_split_parts[::2]  # 取索引0、2、4...的元素，对应每个x分量的训练集
test_x = x_split_parts[1::2]  # 取索引1、3、5...的元素，对应每个x分量的测试集

# 提取y的训练和测试集
y_train, y_test = y_split_parts

验证效果 #

你可以简单检查结构是否匹配：

print(f"原x的元素数量：{len(x)}")
print(f"train_x的元素数量：{len(train_x)}")
print(f"train_x[0]的形状：{train_x[0].shape}（对应原x[0]的训练集）")
print(f"train_x[1]的形状：{train_x[1].shape}（对应原x[1]的训练集）")

更简洁的写法 #

如果喜欢紧凑的代码，可以直接写成：

n = len(x)
train_x, test_x = split_results[0::2][:n], split_results[1::2][:n]
y_train, y_test = split_results[2*n:]

这样就能完美还原和原x结构对应的训练/测试集合啦！

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