Stratified Random Sampling

Stratify=y in the scikit-learn train_test_split function is one argument that people usually include without second thought. This one argument alone is performing stratified sampling. It is something you have done many times… but you don’t know why.

That’s exactly where it should begin. Not a theoretical explanation of what it means. An action you already do, but now named so you can utilize it intentionally.

What Is Stratified Random Sampling?

Consider a dataset used for fraud detection. 100,000 transactions, 95,000 of them being legitimate, while 5,000 are frauds. Simple random sampling would yield 9,720 legitimate and 280 fraud transactions in one attempt – or 9,850 and 150 in another. No guarantee here. In an unlucky attempt, your training data lacks enough fraud examples to recognize the pattern.

This problem can be addressed using stratified random sampling. Split your population into two strata: legitimate transactions and frauds. Sample 9,500 from the former, and 500 from the latter.

Every time.

Stratified random sampling is a probabilistic sampling method that splits a population into several non-overlapping subpopulations – or strata – and then draws a random sample from each, thus making sure all subpopulations are proportionately represented in the resulting dataset.

Two stages are always required: you must stratify first, and then sample. The order matters since sampling first and stratifying afterwards yields you nothing useful – at least, that’s what should be clear in retrospect.

Stratum (singular) refers to one subgroup. Strata (plural) covers all subgroups together. William G. Cochran formalized the statistical theory underpinning this technique in his 1953 textbook Sampling Techniques, though the method had been used informally in survey research for decades before that .

The mechanism that makes this work: members within each stratum are homogeneous relative to your variable of interest; members across strata differ meaningfully. That within-group similarity reduces sampling error. You are not fighting variance from the whole population — you are drawing from smaller, quieter distributions.

Stratified Sampling vs. Simple Random Sampling — The Core Difference

Researchers from a university survey team want to investigate the effects of mental health issues on 10,000 students. International students constitute 8% of the population, totaling 800 individuals. Sample size: 500.

In simple random sampling, 500 samples out of 10,000 would result in 40 international students. However, the actual variance in any single draw would be between 20 and 65 individuals. This level of variance poses an issue, especially for research into international student outcomes.

In stratified sampling, the variance will be fixed at 40 in every draw. Honestly, this is precisely why stratified sampling exists. Not for rigour’s sake but to solve a particular failure case.

The statistical reason: Simple random sampling’s variance is based on the heterogeneity of the total population. On the other hand, stratified sampling uses the sum of variances within each stratum, which are smaller by definition since the strata were formed based on homogeneity. Studies suggest stratified sampling can decrease the standard error in heterogeneous populations by 20 to 40 percent relative to simple random sampling. The maximum reduction occurs when there are many differences between strata and few differences within them.

In a homogenous population, stratified sampling is unnecessary and complex.

Types of Stratified Random Sampling

Proportionate Stratified Sampling

Each stratum contributes to the sample in proportion to its size in the population.

Formula: Sample size of stratum = (Stratum size ÷ Population size) × Total sample size

A university with 6,000 students — 2,400 freshmen (40%), 1,800 sophomores (30%), 1,200 juniors (20%), 600 seniors (10%). You want 200 respondents.

Year	Population	Proportion	Sample
Freshmen	2,400	40%	80
Sophomores	1,800	30%	60
Juniors	1,200	20%	40
Seniors	600	10%	20
Total	6,000	100%	200

Use this when you want your sample to mirror population structure and no single subgroup needs analytical priority.

Disproportionate Stratified Sampling

The sample size allocated to each stratum is purposely not proportional to the stratum size.

A clinical trial investigating rare adverse drug effects: the incidence of adverse events in patients is 2%. Using proportionate allocation sampling, 20 cases of adverse events will be sampled out of 10,000 patients to make up a total of 1,000 sampled subjects. 20 cases will not yield meaningful results. Disproportionate allocation could sample 200 cases of adverse events and 800 controls.

Neyman Optimal Allocation

Most articles stop at proportionate and disproportionate. They shouldn’t.

Neyman allocation — developed by Jerzy Neyman in 1934 — allocates sample size based on both stratum size and stratum variance:

n_h = n × (N_h × σ_h) ÷ Σ(N_i × σ_i)

Strangely, all the guides justify this decision, but do not highlight an important point that if you oversample a stratum, then the population estimates become biased, unless inverse probability weights are used during analysis. Without the use of weights, you produce misleading results. This is no theory; it has been reported in clinical research publications.

where n_h is the number of samples in stratum h, N_h is the population size in stratum h, and σ_h is the standard deviation of stratum h.

The rationale: a high-variance stratum requires more samples to be accurately estimated than a low-variance stratum, which may require only a few samples. Proportional allocation disregards this; Neyman allocation utilizes this knowledge.

Of course, this means that you must know the variances in each stratum beforehand, which is usually not possible. In cases where you have some pilot data or have performed exploratory data analysis on your machine learning dataset, Neyman allocation should be calculated. It will always statistically outperform proportional allocation.

Right. For most practical purposes, proportional allocation is more than enough. Neyman allocation is the optimization step you take after reaching the limits of precision in proportional allocation.

How to Do Stratified Random Sampling — Step by Step

Step 1: Identify the Population and Sampling Frame. The sampling frame consists of the entire exhaustive list of all population members. A partial sampling frame will result in a partial stratum. This is the step that people always rush in the process of doing something quickly. This is also the step which creates all further problems.

Step 2: Identify your stratification variable – choose carefully. While in practice, however, this is often the point at which practitioners will make their first mistake. There are three basic criteria for a good stratification variable: correlation with the dependent variable, creation of homogeneous subgroups, and measurability for all population members.

The wrong stratification variable – one that is not correlated with the dependent variable – only complicates matters. Stratifying a study on blood pressure according to hair color would add nothing to analysis. Stratifying it according to cardiovascular risks, age, and/or medications may alter everything. Test your correlation before making your choice.

Step 3: Confirm that all strata are mutually exclusive and collectively exhaustive. All members of the population should belong to exactly one stratum. There should be no overlap and no exclusion. If an individual aged 45 could belong to both “40-50” and “45-55,” your entire sampling strategy falls apart. Correct the stratum definitions prior to sampling.

Step 4: Determine whether to use proportionate, disproportionate, or Neyman allocation. Use the same criteria as described in the previous section. This decision should be made before handling any data.

Step 5: Determine the number of samples to take from each stratum. Use the proportional allocation formula, or whatever alternative you chose. Be careful about rounding – rounding errors in several strata might affect your overall sample size more than you think.

Step 6: Sample randomly from each stratum. Use a random number generator, a random index selection, or another suitable technique. Convenience sampling has absolutely no place in stratified sampling.

Step 7: Combine and verify. Merge the stratum samples into one dataset. Verify that final class distributions match your target proportions. One line of pandas does this: df[‘label’].value_counts(normalize=True).

Stratified Random Sampling in Python — scikit-learn Implementation

Here’s the thing: every tool you need is already in scikit-learn. Three classes, three different use cases. Know which one to reach for.

Using train_test_split with stratify=

The simplest entry point. The stratify parameter tells scikit-learn to preserve the class distribution of y in both the training and test sets.

python
from sklearn.model_selection import train_test_split
import pandas as pd

# Without stratification -- class balance not preserved
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# With stratification -- class ratios locked in both sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

# Verify the split preserved your class ratios
print("Training set distribution:")
print(y_train.value_counts(normalize=True).round(4))

print("\nTest set distribution:")
print(y_test.value_counts(normalize=True).round(4))

On an imbalanced dataset — 95% class 0, 5% class 1 — without stratify=y, a 20% test split might contain zero class-1 examples by chance. Your model trains on fraud samples and gets tested on none. You would never know until it fails in production.

With stratify=y, both sets maintain the 95/5 split. This is exactly what the parameter is optimizing for.

Using StratifiedShuffleSplit for Multiple Splits

When you need more than one stratified split — for repeated model evaluation or bootstrapping — Stratified Shuffle Split generates multiple independent stratified train-test pairs.

python
from sklearn.model_selection import StratifiedShuffleSplit
import numpy as np

sss = StratifiedShuffleSplit(n_splits=5, test_size=0.2, random_state=42)

for fold, (train_idx, test_idx) in enumerate(sss.split(X, y)):
    X_train, X_test = X[train_idx], X[test_idx]
    y_train, y_test = y[train_idx], y[test_idx]

    unique, counts = np.unique(y_train, return_counts=True)
    dist = dict(zip(unique, counts / len(y_train)))
    print(f"Fold {fold + 1} -- Class distribution: {dist}")

Each of the five splits is independently stratified. A model that performs consistently across five stratified splits is more trustworthy than one evaluated on a single split — because you have tested it against five different random draws, all with the correct class balance.

Stratified K-Fold Cross-Validation

StratifiedKFold is not just a repeated split. It is a full cross-validation procedure where every fold reflects the original class distribution; the training and evaluation sets rotate through the entire dataset.

python
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import accuracy_score
import numpy as np

skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
scores = []

for fold, (train_idx, test_idx) in enumerate(skf.split(X, y)):
    X_train, X_test = X[train_idx], X[test_idx]
    y_train, y_test = y[train_idx], y[test_idx]

    model.fit(X_train, y_train)
    score = accuracy_score(y_test, model.predict(X_test))
    scores.append(score)
    print(f"Fold {fold + 1}: {score:.4f}")

print(f"\nMean accuracy: {np.mean(scores):.4f} ± {np.std(scores):.4f}")

Standard KFold partitions are randomized. With a data set where 3% are positive instances and five folds are used, there will definitely be at least one fold where no positive instances appear in the test partition. The StratifiedKFold avoids that issue by ensuring that each fold has the same 3% proportion of positive instances.

If you are doing imbalanced classification – detecting fraud, diagnosing rare diseases, predicting defects – using StratifiedKFold is non-negotiable.

Worked Examples Across Three Domains

Example 1 — Machine Learning: Fraud Detection Dataset

A fintech company has 100,000 transaction records — 95,000 legitimate (95%) and 5,000 fraudulent (5%). They need 10,000 training samples.

Class	Population	Proportion	Stratified Sample
Legitimate	95,000	95%	9,500
Fraudulent	5,000	5%	500
Total	100,000	100%	10,000

If simple random sampling were used on the same sample, 9,850 genuine and 150 fraudulent samples would be included — a fraud rate of 1.5%, as opposed to 5%. Now, that seems like a difference in favor of stratified sampling, doesn’t it? It is. Why? Because the classification boundary changes;

the precision on the minority class is lower; recall does precisely what it shouldn’t.

Stratified sampling keeps the fraud rate of 5% intact. In every iteration.

Example 2 — Public Health: Vaccine Effectiveness Study

A public health team studies vaccine response across 80,000 city residents. The research question specifically concerns outcomes in the 65+ age group, which is underrepresented proportionately. They apply disproportionate stratified sampling to oversample it.

Age Group	Population	Proportion	Proportionate Sample	Disproportionate Sample
0–17	14,400	18%	180	100
18–44	28,000	35%	350	200
45–64	22,400	28%	280	200
65+	15,200	19%	190	500
Total	80,000	100%	1,000	1,000

The 65+ group receives 500 samples instead of 190. Statistical analysis of elderly outcomes now has adequate power. To be fair, population-level estimates must be weighted back after this — otherwise the study overstates elderly outcomes relative to the full population. That weighting step is where errors are silently introduced in published research.

Example 3 — Market Research: Two-Variable Stratification

A retail company surveys 12,000 customers across four regions and three income brackets. This is nested stratification — two variables, twelve strata simultaneously.

Region / Income	Low	Mid	High	Row Total
North	800	1,200	400	2,400
South	1,500	2,000	500	4,000
East	1,200	1,800	600	3,600
West	600	900	500	2,000
Column Total	4,100	5,900	2,000	12,000

None of the competing articles on this topic show two-variable stratification. But it is common in real market research — region alone misses income-driven preference differences; income alone misses geographic patterns. Combining both creates more precise subgroups, each more internally homogeneous than either variable alone would produce.

The complexity cost is real: twelve strata at a minimum of 20–30 observations each requires 240–360 records before any cell is statistically usable. Count your cells before you commit to the design.

When to Use Stratified Random Sampling — Decision Framework

Use it when:

The population that you have contains distinct subgroups that differ from each other on your dependent variable. The phrase “meaningful difference” implies that the difference is significant for the study, not only statistically significant.

At least one subgroup makes up less than 10% of the population. Under such conditions, simple random sampling becomes quite risky.

You will need to compare these subgroups directly using statistical methods. No inference can be made about any group that was not sampled.

Do not use it when:

There is true homogeneity in the population. The inclusion of strata increases cost with no reduction in variance. Test for this first: If there is little correlation between your stratification variable and the response variable, then stratification does not help.

It is impossible to allocate each individual in the population to one and only one stratum. Strata that overlap or are fuzzy are better than none at all.

To sum up, stratification is worth the trouble only when the groups are truly different and the problem under study depends upon that difference.

Stratified Sampling vs. Other Sampling Methods

Criteria	Simple Random	Systematic	Stratified Random	Cluster
Core mechanism	Every member has equal selection probability	Every nth member selected	Random sampling within predefined strata	Random selection of entire naturally occurring groups
Best for	Homogeneous populations	Ordered, listed populations	Heterogeneous or imbalanced datasets	Geographically dispersed populations
Primary risk	Underrepresents minority groups	Periodicity bias if intervals align with population patterns	Overstratification; strata definition errors	High within-cluster similarity inflates variance estimates
ML use case	Small, balanced datasets	Rarely used	Train-test splits, K-fold cross-validation	Federated learning scenarios
Complexity	Low	Moderate	Moderate	Moderate

Comparison between Quota sampling and Stratified Sampling: Both methods involve dividing a population into sub-groups and then selecting a specific number of individuals from each. The key difference lies in the element of randomness. While Quota sampling involves non-probability sampling, where the researcher selects the individuals to fill the quota, Stratified random sampling involves probability sampling where each individual has a known, non-zero probability of being selected.

Advantages of Stratified Random Sampling

Decreased standard error. Within-stratum variance is always less than the variance of the entire population since you have designed the strata to be homogeneous. Adding up these lower within-stratum variances will yield a lower overall sampling error than simple random sampling in the same population. This decrease is both measurable and meaningful.

Guarantees inclusion of minorities. In cases where a particular group makes up less than 5-10 percent of the population, simple random sampling can be risky. Stratified sampling solves the problem; it actually does that, which is kind of the idea behind it.

Facilitates valid comparison of subgroups. If you do not take a sample specifically designed to include particular subgroups, there is no way you can compare them statistically. Stratification is necessary for making any between-groups inference; without stratification, subgroup comparisons depend on uncontrollable sample makeup.

Cost effectiveness. If your research involves expenses associated with particular subgroups – remote communities, costly laboratory analyses, inaccessible populations – disproportionate allocation maximizes cost-effectiveness while minimizing total sample size.

Increased statistical power for the same sample size. Since variance is decreased, fewer observations will suffice to obtain the same level of confidence and detect effects as in simple random sampling. 

Disadvantages and Common Mistakes

These are indeed problems with stratification: it needs complete information on the population, more planning, and higher complexity of implementation. Overstratification, which leads to an excessive number of strata with insufficient numbers of observations per cell, is indeed a problem and one that frequently occurs.

However, this is not what you will hear from most how-to guides on stratification.

No.

It is not about being “too complex” or “too resource-intensive.” It is wrong in the sense that these errors lead to an incorrect result due to these very errors.

Mistake 1: Stratification with an independent variable. When there is no relationship between the stratification and the outcome variables, stratification increases the complexity of the process but does not decrease its variance. Always check the correlation between the two before proceeding.

Mistake 2: Excessive strata relative to the size of the sample. Twelve strata and a total sample size of 200 will lead to an average of 17 observations per stratum. The reliability of most analyses drops below 20-30 samples per stratum. Count the number of strata you have, calculate the expected sample size per cell, and reconsider your design before collecting any data.

Mistake 3: Not using weights after disproportional sampling. Oversampling some strata means that unless you use inverse probability weights, your estimates about the whole population will be biased. This mistake has been documented in published studies. Clearly document the weights you used and use them consistently.

Mistake 4: Stratifying for ML on variables not present at inference time. If your stratification variable will not be known when running the model in production, stratification is data leakage. Suppose you stratify based on “customer account age” but score customers whose account age is not yet known. In this case, your distribution of training data is different from the distribution at inference time. Think through the entire process of inference before deciding on the stratification variables in ML tasks.

Frequently Asked Questions