Problem 3: Stratified Split for Imbalanced Classes

Problem Description

Write a Python program that performs a stratified train-test split on an imbalanced binary classification dataset, and prints the class distribution in both the training and test sets to confirm proportions were preserved.

Input

A dataset with an imbalanced binary label column.

Output

The class distribution in both the training and test sets.

Constraints

• The dataset should have a clear class imbalance (e.g., 90/10).

Example Input

text
y with 90 zeros and 10 ones

Example Output

text
Train distribution: [72  8]
Test distribution: [18  2]

Concepts Covered

stratify=y, np.bincount().