从鸢尾花到你的数据集：用pandas和sklearn搞定train_test_split的5种真实数据预处理场景

当你第一次在机器学习教程中看到train_test_split时，它通常伴随着整洁的鸢尾花数据集——四个规整的特征列，150条完美无缺的记录。但现实世界的数据更像是一盒被猫玩过的拼图：缺失的碎片、形状各异的板块、还有几张根本不属于这个拼图的卡片。本文将带你跨越从"教科书示例"到"真实项目"的鸿沟，解决那些教程里没告诉你的数据划分难题。

1. 当数据框不完美时：处理缺失值和非数值特征

现实中的数据框很少像iris数据集那样"乖巧"。假设你拿到的是一个电商用户行为数据集：

import pandas as pd
df = pd.DataFrame({
    'user_id': [101, 102, 103, 104, 105],
    'age': [25, 33, None, 45, 28],  # 缺失值
    'gender': ['M', 'F', 'M', None, 'F'],  # 分类特征
    'purchase_amount': [120, 250, 89, 310, 75]
})

处理策略分步指南：

1. 数值型缺失值：用中位数填充比均值更鲁棒

   df['age'] = df['age'].fillna(df['age'].median())
   

2. 分类特征转换：先用众数填充缺失，再进行独热编码

   from sklearn.preprocessing import OneHotEncoder
   df['gender'] = df['gender'].fillna(df['gender'].mode()[0])
   encoder = OneHotEncoder(sparse=False)
   gender_encoded = encoder.fit_transform(df[['gender']])
   

3. 最终划分方案：

   from sklearn.model_selection import train_test_split
   X = pd.concat([df[['user_id', 'age', 'purchase_amount']], 
                pd.DataFrame(gender_encoded)], axis=1)
   y = df['purchase_amount'] > 200  # 假设我们要预测大额消费
   X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3)
   

注意：永远先处理缺失值再进行数据划分，避免信息泄露

2. 类别不平衡时的生存指南：stratify参数的艺术

在信用卡欺诈检测中，正样本可能只占0.1%。直接随机划分会导致测试集可能完全没有正样本。这时stratify参数就是你的救星：

# 模拟一个严重不平衡的数据集
import numpy as np
X = np.random.rand(10000, 10)  # 10000个样本，10个特征
y = np.array([0]*9900 + [1]*100)  # 1%的正样本

# 错误做法：直接划分
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
print("Test set positive ratio:", y_test.mean())  # 可能为0

# 正确做法：分层抽样
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, stratify=y)
print("Test set positive ratio:", y_test.mean())  # 保持1%

分层抽样的进阶技巧：

• 当目标变量是连续值时，可以先分箱再分层
• 多标签分类时，可使用sklearn.multiclass中的策略
• 对于极度不平衡数据(如1:1000)，考虑先过采样再分层

3. 时间序列数据的陷阱：为什么必须禁用shuffle

处理股票价格或传感器数据时，时间顺序就是生命线。一个常见的错误示范：

# 错误的时间序列划分
dates = pd.date_range('2023-01-01', periods=100)
stock_prices = np.cumsum(np.random.randn(100))  # 随机游走模拟股价
X = stock_prices.reshape(-1, 1)
y = np.roll(stock_prices, -1)  # 用今日价格预测明日

# 绝对不要这样做！
X_train, X_test, y_train, y_test = train_test_split(X, y, shuffle=True)

正确的时间序列划分应该是：

# 正确的时间序列划分
test_size = 0.2
split_point = int(len(X) * (1 - test_size))
X_train, X_test = X[:split_point], X[split_point:]
y_train, y_test = y[:split_point], y[split_point:]

# 或者使用TimeSeriesSplit
from sklearn.model_selection import TimeSeriesSplit
tscv = TimeSeriesSplit(n_splits=5)
for train_index, test_index in tscv.split(X):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]

时间序列验证的关键原则：

• 测试集时间必须晚于训练集
• 考虑季节性因素时，保持完整周期不被分割
• 对于滚动预测任务，使用时间窗口交叉验证

4. 数据结构大杂烩：统一处理不同格式的输入

真实项目中，特征可能分散在多个数据结构中。比如：

# 不同格式的特征数据
feature_df = pd.DataFrame({'age': [25, 30, 35], 'income': [5000, 8000, 6000]})
feature_list = [[1, 0], [0, 1], [1, 1]]  # 来自其他系统的特征
feature_array = np.random.rand(3, 3)  # 图像提取的特征

# 解决方案1：先拼接再划分
from sklearn.preprocessing import StandardScaler
X_concat = np.concatenate([
    feature_df.values,
    np.array(feature_list),
    feature_array
], axis=1)
X_train, X_test = train_test_split(X_concat, test_size=0.3)

# 解决方案2：分别划分再组合（保持对应关系）
indices = np.arange(len(feature_df))
train_idx, test_idx = train_test_split(indices, test_size=0.3)
X_train = {
    'tabular': feature_df.iloc[train_idx],
    'list_feat': [feature_list[i] for i in train_idx],
    'array_feat': feature_array[train_idx]
}

多源数据划分的黄金法则：

1. 确保所有特征矩阵的样本顺序一致
2. 对于非数值数据，先转换或保持索引对应关系
3. 考虑使用pandas.Index来维护样本标识

5. 划分后的健康检查：验证分布一致性的5种方法

数据划分后不做分布检查，就像做完手术不缝合——迟早出问题。以下是必备检查清单：

数值特征检查：

# 使用KS检验比较分布
from scipy.stats import ks_2samp
for col in X_train.columns:
    stat, p = ks_2samp(X_train[col], X_test[col])
    print(f"{col}: p-value={p:.3f}")  # p>0.05表示分布一致

分类特征检查表：

检查项	方法	可接受标准
类别比例	value_counts(normalize=True)	差异<5%
新类别	检查测试集独有类别	应当为0
稀有类别	最小类别样本数	测试集≥3个

标签分布可视化：

import matplotlib.pyplot as plt
plt.figure(figsize=(10, 4))
plt.subplot(121)
y_train.value_counts().plot(kind='bar', title='Train')
plt.subplot(122)
y_test.value_counts().plot(kind='bar', title='Test')
plt.show()

高级检查技巧：

• 使用pandas_profiling对比报告
• 对高维数据做PCA后比较分布
• 训练一个分类器区分训练/测试集（AUC应接近0.5）

在实际项目中，我发现最常被忽视的是分类特征的稀有类别问题。曾经在一个客户分群项目中，测试集出现了训练集没有的职业类别，导致模型线上表现灾难性下降。现在我的团队强制要求对所有分类特征执行以下检查：

# 分类特征完整性检查
train_categories = set(X_train['category_column'].unique())
test_categories = set(X_test['category_column'].unique())
assert test_categories.issubset(train_categories), "发现新类别！"