Scikit-learn: train/test split to include have same representation of two different types of values in a column

Question

I have a dataset of online purchase orders that contains two types of customers:

• Customers who have an account and thus are known customers with unique customer number.
• Customers who do not have an account and do their purchases as a guest; thus, these are unknown customers and are given a default number(out of a list of 50 pre-defined numbers). Meaning it is not a unique number and the only way to distinguish transactions here is by the unique order number.

To avoid distortion and make sure that my train/test dataset will have the same distribution when it comes to known and unknown customers, I want to select a Train/Test set as follows:

The part of the dataset that includes the known customers to be split 80/20.

The part of the dataset that includes the unknown customers to be split 80/20 ( the selection of the sub dataset here will depend on the rule that orders are not linked to a customer (= linked to a default customer number).

In Scikit-learn I need to set up the X(features) and Y(target) in order to do the train_test_split. My question is would it be logical and doable to split my dataframe into:

• df1 which contains all known customers.
• df2 which contains all unknown customers.

Then set up X, Y and consequently x_train, x_test, y_train , y_test for each of df1 and df2, then merge the result into one and continue to build the rest of the model?

Is this a good approach or is there another solution?

Answer 1

Keep only one dataframe
Add a column(if not available) to mark - Guest Or Customer
Then, simply split with stratify flag on that column

from sklearn.model_selection import train_test_split
x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=19, stratify=data['guest'])

stratify : array-like, default = None
If not None, data is split in a stratified fashion, using this as the class labels

Answer 2

I agree to opinions said before. Just as alternative, if you see that customer behavior is too different if it is a guest or not, depending also on model you use, probably it would make sense to use two different models. For example, if you know will use LogisticRegression and not regular customers behavior is distributed in bigger range, then probably you achieve better scores by using two Logistic Regressions (one for regular customers, another for guest customers).

Answer 3

Welcome to Data Science at StackExchange,

One way to accomplish this is to use the stratify option in train_test_split, since you are already using that function (this will also work for ensuring your labels are equally distributed, very useful in modelling an unbalanced dataset):

Train,Test = train_test_split(df, test_size=0.50, stratify=df['B'])

In my example, you can see that there are 2 values in column B, and they are equally distributed between the 2 datasets. In your case, B would be the column indicating those customers with an account and those without.

Answer 4

Would it be a possibility to do the following:

keep one dataset but give those unknown customers a unique number per unique order number.

Something like updating the customer code column with the same code like the unique order number prefixed with something that indicates that it was an unknown customer before.