How to split data into train and test sets with scikit-learn

Model evaluation is meaningful only when test rows stay outside model fitting and learned preprocessing. A held-out partition creates that boundary, while scikit-learn keeps each feature row aligned with its label when the dataset is divided.

The train_test_split() helper accepts multiple arrays, sparse matrices, or dataframes with the same row count and returns a training and test portion for each input. An integer random_state makes the selected rows repeatable for unchanged inputs, while stratify=y keeps the class proportions close across a classification split.

A random row split assumes samples can be shuffled independently. Grouped subjects, repeated measurements, or time-ordered observations need a group-aware or time-series splitter because related or future information could otherwise cross the evaluation boundary.

1. Create split_train_test.py with the imports and aligned input arrays.

   split_train_test.py

   from collections import Counter
    
   import numpy as np
   import sklearn
   from sklearn.datasets import load_breast_cancer
   from sklearn.model_selection import train_test_split
    
    
   dataset = load_breast_cancer()
   X = dataset.data
   y = dataset.target
   row_ids = np.arange(X.shape[0])

   The row_ids array travels through the same split as X and y, which makes row overlap directly measurable without changing the feature data.

2. Append the stratified train/test split with a 25 percent holdout.

   X_train, X_test, y_train, y_test, train_ids, test_ids = train_test_split(
       X,
       y,
       row_ids,
       test_size=0.25,
       stratify=y,
       random_state=42,
   )

   test_size=0.25 reserves one quarter of the rows for testing. stratify=y keeps both target classes represented, and the integer random_state fixes the shuffled row selection for unchanged inputs.

3. Append a repeated split that selects the test row IDs again.

   _, _, _, _, _, repeated_test_ids = train_test_split(
       X,
       y,
       row_ids,
       test_size=0.25,
       stratify=y,
       random_state=42,
   )

   The underscores discard the repeated feature, label, and training arrays because only the second set of test row IDs is needed for the reproducibility check.

4. Append a helper that converts numeric label counts to dataset class names.

   def class_counts(labels):
       counts = Counter(labels)
       return {
           str(dataset.target_names[int(label)]): int(count)
           for label, count in sorted(counts.items())
       }

5. Append the shape, class-balance, overlap, and reproducibility checks.

   print(f"scikit-learn {sklearn.__version__}")
   print(f"total rows: {X.shape[0]}")
   print(f"train shape: {X_train.shape}")
   print(f"test shape: {X_test.shape}")
   print(f"train class counts: {class_counts(y_train)}")
   print(f"test class counts: {class_counts(y_test)}")
   print(f"test share: {len(test_ids) / len(row_ids):.3f}")
   print(f"overlap rows: {len(np.intersect1d(train_ids, test_ids))}")
   print(f"repeat split matches: {np.array_equal(test_ids, repeated_test_ids)}")

6. Run the completed split script.

   $ python3 split_train_test.py
   scikit-learn 1.9.0
   total rows: 569
   train shape: (426, 30)
   test shape: (143, 30)
   train class counts: {'malignant': 159, 'benign': 267}
   test class counts: {'malignant': 53, 'benign': 90}
   test share: 0.251
   overlap rows: 0
   repeat split matches: True

7. Verify that the output proves a separate, stratified, repeatable holdout.

   The 569 rows become 426 training rows and 143 test rows, and both subsets contain malignant and benign labels. overlap rows: 0 proves no row ID belongs to both subsets, while repeat split matches: True proves the same integer random state selects the same held-out rows.