Stratified k-fold assignment
Medium12 min11 hidden tests
A fraud dataset has far fewer fraud cases than normal ones. If you cut it into k folds for cross-validation by position, one fold can end up with no fraud at all, and its score means little. Write stratified_folds(labels, k) that gives each example a fold number so every class is spread evenly across the folds. Deal the examples out like cards: line them up class by class, then hand them to folds 0, 1, ..., k - 1, 0, 1, ... in that order.
Tap an example to flip its class
line up by class, deal to folds 0, 1, …, k − 1, 0, …examples, in their original order (tap to flip)
lined up by class, then dealt
folds
fold 0 · 4 examples, 1 fraud
fold 1 · 3 examples, 1 fraud
fold 2 · 3 examples, 1 fraud
Lined up by class, normal first: 0, 1, 3, 4, 5, 7, 8, 2, 6, 9. Dealt to folds 0, 1, 2, … without restarting at the fraud cases, so they land in folds 1, 2, 0. The answer, by example: [0, 1, 1, 2, 0, 1, 2, 2, 0, 0].
Folds 0, 1, 1, 2, 0, 1, 2, 2, 0, 0.
Examples
Four normal cases and two fraud cases, split evenly
stratified_folds([0, 1, 0, 0, 1, 0], 2)
returns [0, 0, 1, 0, 1, 1]
Three classes, one of each per fold
stratified_folds([2, 0, 1, 0, 2, 1, 0, 1, 2], 3)
returns [0, 0, 0, 1, 1, 1, 2, 2, 2]
Requirements
labelsis a list or 1-D array of n int class labels.kis an int from 2 to n.- Return a 1-D NumPy int array of length n, where entry i is the fold (0 to k - 1) of example i.
- Line the examples up by class, smallest label first. Inside a class, keep their original order.
- Deal that line to folds 0, 1, ..., k - 1 and start over at 0, without restarting when a new class begins.
The idea behind it
Cross-validation splits a dataset into k folds, trains on k - 1 of them and tests on the one left out, k times. The average score is more reliable than a single split. When one class is rare, a plain split can give a fold with almost none of it. That fold's recall is then noise, or undefined. Stratified k-fold fixes this. It keeps each class's share in every fold close to its share in the whole dataset. Dealing the examples like cards, class by class, does exactly that: a class with m examples puts either floor(m / k) or ceil(m / k) of them in each fold. Continuing the deal across classes also keeps the fold sizes within one of each other. Libraries add a shuffle first. Here the order is fixed so the result is easy to check.