如何用Pandas原生方法无循环按customer_id拆分训练测试集？

当然可以用Pandas原生的矢量化操作来实现，完全不需要写循环，这样代码不仅更简洁，运行效率也会比循环高很多，尤其是当你的数据集比较大的时候。

优化后的完整代码 #

import pandas as pd
import numpy as np
import random

# 读取并预处理数据
df_train = pd.read_csv('data/train.csv')
print('原始df_train.shape', df_train.shape)
df_train = df_train.replace(np.nan, 'nan', regex=True)

# 获取所有唯一的customer_id
train_customer_id_set = df_train.customer_id.unique()
print('唯一customer_id数量', len(train_customer_id_set))

# 随机选择测试用的customer_id
n = 1000
test_customer_id_set = list(train_customer_id_set)
random.shuffle(test_customer_id_set)
test_customer_id_set = test_customer_id_set[:n]

# --- 核心优化部分：替代循环 ---
# 构建测试集：直接用isin筛选
df_test = df_train[df_train['customer_id'].isin(test_customer_id_set)]

# 从训练集中移除测试集数据：用~取反筛选
df_train = df_train[~df_train['customer_id'].isin(test_customer_id_set)]
# -----------------------------

train_customer_id_set = df_train.customer_id.unique()
print('处理后df_train.shape', df_train.shape)
print('df_test.shape', df_test.shape)

关键优化点解释 #

• 构建df_test：使用df_train['customer_id'].isin(test_customer_id_set)生成布尔索引，直接筛选出所有属于测试用户的行，这是Pandas的矢量化操作，比循环逐个拼接DataFrame高效得多。
• 更新df_train：用~对布尔索引取反，筛选出不属于测试用户的行，直接得到新的训练集，避免了循环调用drop的繁琐操作，同时性能更好。

这种方法完全贴合Pandas的设计理念，利用其内置的矢量化运算能力，既简化了代码逻辑，又提升了运行效率。

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