如何合并平衡后的X_train与y_train并保留plco_id？

一、合并X_train与y_train的正确方式 #

SMOTEENN的fit_resample方法会严格保持X和y的行对应关系——生成的新样本和原样本的索引顺序完全对齐，直接用pd.concat按列合并即可：

import pandas as pd

# 合并X_train和y_train
train_df = pd.concat([X_train, y_train], axis=1)

验证匹配正确性的小技巧 #

可以通过以下方式确认合并后的数据没有错位：

# 检查三者行数是否一致
assert len(X_train) == len(y_train) == len(train_df)

# 随机抽取一行验证y值匹配
random_idx = 100
assert train_df.loc[random_idx, 'lung_cancer'] == y_train.loc[random_idx, 'lung_cancer']

二、在拆分与平衡过程中保留plco_id的操作方法 #

要保留ID，需要调整初始的X和y定义，让plco_id随X一起参与后续流程，具体步骤如下：

1. 调整原始数据拆分逻辑 #

把plco_id留在X中，只移除目标变量lung_cancer：

# 仅移除目标变量，保留plco_id在X里
X = df2.drop('lung_cancer', axis=1)
y = df2['lung_cancer']

2. 执行分层拆分（无需修改原有代码） #

原有的train_test_split代码可以直接用，因为X现在包含plco_id，拆分后X_temp、X_train等数据集都会自动保留该列。

3. 平衡训练集时保留plco_id #

plco_id是标识列，不需要参与过采样/欠采样，这里提供两种可行方式：

方式一：先分离ID列，平衡后再合并

# 从训练集中提取ID列
train_ids = X_train['plco_id']
# 移除ID列，准备做平衡处理
X_train_features = X_train.drop('plco_id', axis=1)

# 执行SMOTEENN平衡
smt = SMOTEENN(random_state=122)
X_train_balanced, y_train_balanced = smt.fit_resample(X_train_features, y_train)

# 重新合并ID列（过采样生成的新样本没有真实ID，可用NaN填充或自定义规则）
X_train_balanced = pd.concat([
    train_ids.reindex(X_train_balanced.index).reset_index(drop=True),
    X_train_balanced.reset_index(drop=True)
], axis=1)

# 最终合并成包含ID的训练DataFrame
train_df_with_id = pd.concat([X_train_balanced, y_train_balanced.reset_index(drop=True)], axis=1)

方式二：用ColumnTransformer让SMOTEENN忽略ID列（推荐）

通过转换器指定仅对特征列做过采样，ID列直接保留：

from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import FunctionTransformer

# 定义需要参与过采样的特征列（排除plco_id）
feature_cols = [col for col in X_train.columns if col != 'plco_id']

# 创建转换器：特征列用SMOTEENN处理，ID列直接传递
preprocessor = ColumnTransformer(
    transformers=[
        ('smoteenn', smt, feature_cols),
        ('passthrough', FunctionTransformer(), ['plco_id'])
    ])

# 拟合并转换训练集
X_train_balanced = preprocessor.fit_resample(X_train, y_train)
# 转换为DataFrame（fit_resample返回数组，需恢复列名）
X_train_balanced = pd.DataFrame(X_train_balanced, columns=feature_cols + ['plco_id'])
y_train_balanced = pd.Series(y_train_balanced, name='lung_cancer')

# 合并为最终的带ID训练DataFrame
train_df_with_id = pd.concat([X_train_balanced, y_train_balanced], axis=1)

注意事项 #

• 过采样生成的新样本没有真实plco_id，可根据需求用NaN、重复现有ID或自定义标识（如new_sample_123）填充。
• 验证集和测试集不需要平衡，它们的plco_id会完整保留原始数据的标识。

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