📊 Understanding train_test_split in Scikit-learn: How and Why to Split Your Data for Machine Learning

When building a machine learning model, one of the first and most important steps is splitting your dataset. The most commonly used function for this is train_test_split, which comes from scikit-learn, one of Python’s most powerful machine learning libraries.

In this article, we’ll explore:

• What train_test_split does

• How to use it in practice with real examples

• A common question about category mismatches between training and test sets

• And the importance of a validation set

🔍 What Is train_test_split?

The train_test_split() function is used to divide your dataset into training, validation, and testing subsets (or just training and testing in simpler setups). This ensures that:

• The model is trained on one portion of the data

• Tuned and evaluated on a separate validation set

• Finally tested on unseen data

This practice answers the key question:

“How well will my model perform on data it has never seen before?”

✅ Basic Syntax

pythonCopyEditfrom sklearn.model_selection import train_test_split

# Split data: 80% for training, 20% for testing
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42 , stratify=y)

• X = your features (input variables), like "CustomerID", "Contract", "MonthlyCharges"

• y = your label (target), like "Churn"

• test_size=0.2 → 20% of the data will be used for testing

• random_state=42 → ensures reproducible splits

• stratify=y → ensures the train/test sets have similar churn ratios

🧪 Real-World Example: Splitting Customer Churn Data (with Validation)

Here’s a realistic dataset of customer churn:

from sklearn.model_selection import train_test_split
import pandas as pd

# Sample dataset
data = {
    "CustomerID": [101, 102, 103, 104, 105, 106, 107, 108, 109],
    "Contract": ["Monthly", "Two year", "Monthly", "One year", "Monthly", "Two year", "One year", "Monthly", "Two year"],
    "MonthlyCharges": [29.9, 56.5, 42.3, 70.1, 99.9, 55.5, 65.0, 33.1, 89.4],
    "Churn": [1, 0, 1, 0, 1, 0, 0, 1, 0]
}


# Features and target
# X is a new dataframe with two columns: "Contract" and "MonthlyCharges" from the dataset.
# Y is actually a series since it has only 1 column It selects the "Churn" column.
# Most models (like in scikit-learn) expect y to be a Series (not a DataFrame), 
# because they only need a single target variable.
# These are the input variables the model will use to learn patterns and make predictions.
df = pd.DataFrame(data)
X = df[["Contract", "MonthlyCharges"]]
y = df["Churn"]

# 60% Train, 20% Validation, 20% Test
X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.25, random_state=42, stratify=y_temp)

print("X_train:\n", X_train)
print("\nX_val:\n", X_val)
print("\nX_test:\n", X_test)

result

X_train:
   Contract  MonthlyCharges
0   Monthly            29.9
7   Monthly            33.1
2   Monthly            42.3
6  One year            65.0
1  Two year            56.5

X_val:
   Contract  MonthlyCharges
3  One year            70.1
4   Monthly            99.9

X_test:
   Contract  MonthlyCharges
8  Two year            89.4
5  Two year            55.5

📌 Important Question: What if a Category in Training Set Is Missing from Test Set?

In this case, let’s say X_train contains "One year" under the Contract column, but X_test does not.

Here’s what happens:

You might ask:

❓ “Will it be a problem if X_test doesn’t contain all the categories from X_train?”

✅ Answer:

No — this is not a problem at all.

The model has already seen and learned the "One year" category during training, so it’s prepared to make predictions even if that category doesn’t appear in the test set.

✅ This is safe and expected behavior.

🚨 But What If the Test Set Has a New Category?

Let’s flip the situation.

Imagine that "Three year" appears only in the test set, not in the training set. Then, during prediction, if you’re using a OneHotEncoder without precautions, this will cause an error:

plaintextCopyEditValueError: Found unknown categories ['Three year'] in column 0 during transform

✅ How to Prevent This:

Use this when setting up your encoder:

pythonCopyEditfrom sklearn.preprocessing import OneHotEncoder

encoder = OneHotEncoder(handle_unknown="ignore")

This way, the encoder ignores unseen categories and safely continues.

🛠 Recap: Best Practices for train_test_split

---

📝 Final Thoughts

The train_test_split function is simple but essential. It’s the foundation for evaluating any machine learning model fairly. Understanding how it works — especially with categorical data — will save you hours of debugging and ensure your results are trustworthy.

Next time you’re preparing your data, remember: training on one set and testing on another isn’t just optional — it’s a must.

🧠 Conceptual Questions

1. Why do we split data into training and testing sets?
→ To evaluate how well the model generalizes to unseen data.

2. What does the test_size parameter control?
→ The percentage or number of samples used for testing.

3. What happens if we don't use a random_state?
→ Each run may produce different splits, leading to inconsistent results.

4. Why is it bad to train and test a model on the same data?
→ It causes overfitting—the model memorizes the data instead of learning patterns.

5. What's the difference between test_size=0.2 and train_size=0.8?
→ Both give the same result if they add up to 1.0—just different ways of expressing the split.

💻 Code-Oriented Questions

6. How do I split a dataset into training, validation, and test sets?
→ First split into temp + test, then split temp into train + validation.

X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.2)
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.25)  
# 0.25 * 0.8 = 0.2

7. How do I use stratify=y and what does it do?
→ Maintains the same class distribution in both train and test sets.

train_test_split(X, y, test_size=0.2, stratify=y)

8. Can I use train_test_split with NumPy arrays instead of Pandas?
→ ✅ Yes, it works with lists, arrays, or DataFrames.

9. How can I shuffle data before splitting?
→ It's shuffled by default. Use shuffle=False to turn it off.

10. Can I apply train_test_split on multiple arrays (e.g., X, y, sample weights)?
→ ✅ Yes, pass all arrays in order:

pythonCopyEditX_train, X_test, y_train, y_test, w_train, w_test = train_test_split(X, y, w, test_size=0.2)

---

🌍 Real-World Scenario Questions

11. What should I do if my test set is too small to be representative?
→ Increase test_size (e.g., use 30%) or use cross-validation.

12. How do I handle unseen categories in the test set during encoding?
→ Use OneHotEncoder(handle_unknown="ignore").

13. How should I split time series data instead of using train_test_split?
→ Use chronological split, not random—consider TimeSeriesSplit.

14. When should I use cross-validation instead of train_test_split?
→ Use cross-validation for more robust performance estimation, especially on small datasets.