stratified_split

Builtin index · Statistics & ML

result = stratified_split(X, y, [test_size=0.2], [seed=])

Input: X, y, test_size, seed · Output: result

What it is

The same shape as train_test_split on X/y/test_size/seed=, but splits each class of y separately before combining, so a rare class keeps its proportion in both sets every time rather than only on average.

Overloads

One call shape — this function is not overloaded.

Examples

From the documentation, and run by the build -- the output and the figure below are what the engine actually produced. Simplest first.

seq = sequential_split(X, y, test_size = 0.2)
print("sequential_split train rows: {rows(seq.X_train)}")

strat = stratified_split(Xc, yc, test_size = 0.3, seed = 9)
print("stratified_split train rows: {rows(strat.X_train)}")

tvt = train_val_test_split(Xc, yc, val_size = 0.2, test_size = 0.2, seed = 9)
print("train/val/test rows: {rows(tvt.X_train)}, {rows(tvt.X_val)}, {rows(tvt.X_test)}")

m = knn_model(strat.X_train, strat.y_train, 5, kind = "classification")
confusion_matrix(strat.y_test, m.predict(strat.X_test), 2)

sequential_split train rows: 96
stratified_split train rows: 140
train/val/test rows: 120, 40, 40

from statistics-ml

How it works

stratified_split(X, y, [test_size=0.2], [seed=]) — scikit-learn's train_test_split(..., stratify=y): rows are first grouped by their exact y value (first-seen group order), each group is independently shuffled (one seeded RNG stream shared across all groups, so seed= still reproduces the whole split byte-for-byte) and split at the same test_size fraction, then every group's train/test rows are concatenated. The point (unlike a plain random split, which only preserves class proportions *on average*): a rare class is split into the same proportions in both train and test every time, not just "some points from each class landed somewhere" — worth checking with real numbers, not just "didn't crash", on any dataset with unequal class sizes.

From the note above this builtin in the interpreter, where the reasoning was written down as it was built.