Train/Test Split & Shuffling in C++

June 2026 · 10 min read

Every model in this series ends with a line like "split 80/20, train, and report accuracy." This post implements that split. The basic version is a shuffle and a cut; the better version, for imbalanced data, is a stratified split that keeps each class's proportion intact in both halves.

Why Shuffle First

Data often arrives sorted, by class, by time, by collection batch. Cutting the last $20\%$ off a class-sorted file could put an entire class in the test set and none in training. Shuffling first breaks any such ordering so the split is representative. The plain split shuffles the row indices and cuts:

#include <bits/stdc++.h>
using namespace std;

// plain shuffle-and-cut split: returns row indices for train and test
void trainTestSplit(const vector<int>& y, double testFrac, unsigned seed,
                    vector<int>& trainIdx, vector<int>& testIdx) {
    int N = y.size();
    vector<int> idx(N);
    iota(idx.begin(), idx.end(), 0);
    shuffle(idx.begin(), idx.end(), mt19937(seed));
    int nTest = round(testFrac * N);
    testIdx.assign(idx.begin(), idx.begin() + nTest);
    trainIdx.assign(idx.begin() + nTest, idx.end());
}

Splitting on indices rather than copying rows is the same trick the decision tree uses: the feature matrix is never duplicated, the split just hands each side a list of row numbers.

The Problem with Imbalance

Suppose only $30\%$ of the rows are positive. A random split can drift: by chance the test set might be $35\%$ positive and the train set $29\%$. On a small or heavily skewed dataset that drift distorts the score. A stratified split fixes it by splitting each class separately and then merging, so every class contributes the same fraction to test in proportion to its size.

// stratified split: split each class separately so ratios are preserved
void stratifiedSplit(const vector<int>& y, double testFrac, unsigned seed,
                     vector<int>& trainIdx, vector<int>& testIdx) {
    map<int, vector<int>> byClass;
    for (int i = 0; i < (int)y.size(); ++i) byClass[y[i]].push_back(i);

    mt19937 rng(seed);
    for (auto& [cls, rows] : byClass) {
        shuffle(rows.begin(), rows.end(), rng);
        int nTest = round(testFrac * rows.size());
        for (int k = 0; k < (int)rows.size(); ++k)
            (k < nTest ? testIdx : trainIdx).push_back(rows[k]);
    }
}

Running It

Build an imbalanced label vector that is about $30\%$ positive, then compare the two splits' positive rates:

int main() {
    mt19937 rng(0);
    vector<int> y;
    for (int i = 0; i < 1000; ++i)
        y.push_back(generate_canonical<double,10>(rng) < 0.30 ? 1 : 0);

    auto posRate = [&](const vector<int>& idx) {
        int pos = 0; for (int i : idx) pos += y[i];
        return (double)pos / idx.size();
    };
    auto report = [&](const char* tag, const vector<int>& tr, const vector<int>& te) {
        printf("%-14s -> train %.3f  test %.3f\n", tag, posRate(tr), posRate(te));
    };

    vector<int> tr, te;
    trainTestSplit(y, 0.2, 1, tr, te);     report("random split",  tr, te);
    tr.clear(); te.clear();
    stratifiedSplit(y, 0.2, 1, tr, te);    report("stratified",    tr, te);
}

The random split's test set drifts to $0.315$ positive while training sits at $0.291$, a gap of more than two points away from the true $0.296$. The stratified split pins both halves to $\approx 0.296$, matching the population. On balanced, plentiful data the plain split is fine, which is why the model posts use it; reach for stratification when a class is rare or the dataset is small.