我想根据customer_id
(数据框中的多行可能有相同的customer_id
)将训练数据拆分为训练/测试集,我在想是否有更符合Pandas原生方式的方法来完成build df_test
和drop from df_train
部分,而不使用循环?
# 将数据拆分为训练/测试集
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)
train_customer_id_set = df_train.customer_id.unique()
print('len(train_customer_id_set)', 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]
# 问题:如何在不使用循环的情况下完成?
# 构建df_test
df_list = []
for customer_id in test_customer_id_set:
df = df_train[df_train['customer_id']==customer_id]
df_list.append(df)
df_test = pd.concat(df_list)
# 从df_train中删除
for customer_id in test_customer_id_set:
df_train = df_train.drop(df_train[df_train.customer_id==customer_id].index)
train_customer_id_set = df_train.customer_id.unique()
print('df_train.shape', df_train.shape)
print('df_test.shape', df_test.shape)
回答:
在计算test_customer_id_set
之后,你所做的事情似乎等同于:
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)]