CODEWITHSUNDEEP
pythonpandasmachine-learning
mrgloom
mrgloom

Reputation: 21622

Pandas: get subset of data without cycle

I'm trying to split train data to train/test split based on customer_id (several rows in dataframe can have same customer_id) and I wonder can we do build df_test and drop from df_train sections without a cycle in more pandas-native way?

#Split data for train / test split

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))

#Split train data to train/test by 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]

#Q: how to do it without cycle?

#build 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)

#drop from 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)

Upvotes: 1

Views: 83

Answers (1)

Ami Tavory
Ami Tavory

Reputation: 76297

Following the point where you calculated test_customer_id_set, it seems like what you're doing is equivalent to:

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)]

Upvotes: 2

Related Questions