How to Perform Stratified Sampling in Pandas

Use Pandas to sample within groups defined by a stratum column. Choose an equal number from every group when group comparisons matter, or allocate a fixed total in proportion to group sizes when the sample should reflect the population mix. A fixed random seed makes an illustrative draw reproducible.

Stratified Sampling in Statistics

In stratified sampling, divide a population into mutually exclusive groups called strata, then select observations from each group. A stratum might be a grade, region, age band, or another characteristic relevant to the question. Sampling within these groups can ensure that important groups appear in a sample; it does not by itself remove selection bias or guarantee that the sample represents the population. The sampling frame and the way observations are chosen still matter.

Researchers use stratification when groups differ in ways relevant to an estimate, and data practitioners may use it when creating train and test sets so that a category is represented in both. Decide which column defines the strata and whether each group should contribute the same count or a share based on its population size. The two choices answer different questions: equal counts emphasize comparisons among groups, while proportional counts preserve their relative sizes.

Perform Stratified Sampling in Pandas

The example below creates a student DataFrame with ten rows in three grade groups. It takes two students per grade for an equal-sized sample, then allocates a total of six students proportionally. The proportional allocation uses the largest-remainder method: take the floor of each group’s exact quota, then assign remaining places to the groups with the largest fractional remainders. Here the group sizes are 5, 3, and 2, so the six-row allocation is 3, 2, and 1.

import pandas as pd

students = {
    "Name": [
        "sanay", "shivesh", "rutwik", "preet", "yash",
        "mann", "pritesh", "hritesh", "raj", "tarun",
    ],
    "ID": ["001", "002", "003", "004", "005", "006", "007", "008", "009", "010"],
    "Grade": ["A", "A", "C", "B", "B", "B", "C", "A", "A", "A"],
    "Category": [2, 3, 1, 3, 2, 3, 3, 1, 2, 1],
}
df = pd.DataFrame(students)

# Take the same number from each grade.
equal_sample = df.groupby("Grade", group_keys=False).sample(n=2, random_state=7)
print("Equal-sized counts:")
print(equal_sample["Grade"].value_counts().sort_index().to_dict())

# Allocate a fixed total in proportion to each grade's population size.
target = 6
group_sizes = df["Grade"].value_counts().sort_index()
quotas = group_sizes / len(df) * target
allocation = quotas.astype(int)
places_left = target - int(allocation.sum())
remainders = (quotas - allocation).sort_values(ascending=False)
for grade in remainders.index[:places_left]:
    allocation.loc[grade] += 1

proportional_sample = pd.concat(
    [
        group.sample(n=int(allocation.loc[grade]), random_state=7)
        for grade, group in df.groupby("Grade", sort=True)
    ],
    ignore_index=True,
)
print("Proportional counts:")
print(proportional_sample["Grade"].value_counts().sort_index().to_dict())

DataFrame.groupby() forms the groups used for sampling. See the localized Pandas groupby() reference for more details.

For an equal-sized sample, sample(n=2) requests two rows from every grade. The smallest group therefore needs at least two rows; choose a smaller per-group count or revise the strata if it does not. For a proportionate sample, target is the requested total and the allocation step handles rounding so the group counts sum to that total. Different rounding rules can assign a remaining place to a different group when fractional remainders tie, so use a documented rule for repeatable analysis.

Both calls use random_state=7 so the same input and Pandas environment produce a repeatable draw. Change the seed to obtain another random draw. Stratification controls how many observations come from each group; it does not make observations random within a group if the source data or sampling process is biased. Report the stratum column, allocation rule, target size, and seed when the sample needs to be reproduced.