CreateStratifiedRandomSampleGroups

Introduction

Simple random assignment can accidentally put more of one demographic or class into one group than another, especially with small samples — undermining the comparison you're trying to make. CreateStratifiedRandomSampleGroups answers "how do I split this population into groups that are not just random, but provably balanced on the categories that matter?" by stratifying the random assignment so key categorical proportions hold across every group. Reach for it for A/B tests, clinical trials, or cross-validation folds where class or demographic balance across groups is a requirement, not a nice-to-have.

Teaching Note

The function is particularly useful for:

Designing A/B tests where groups must be balanced by demographics (e.g., age or region)
Creating cross-validation folds that preserve class balance for machine learning
Designing clinical or social research trials where treatment groups must be comparable
Splitting marketing segments to ensure representative testing across customer tiers
Educational research where student groups must be balanced by prior performance
Validating model performance across different sub-populations in a balanced way

Parameters

Parameter	Default	Description
`dataframe`required	—	The input pandas DataFrame containing the records to be partitioned.
`number_of_groups`required	—	The number of discrete groups to create. Records within each stratum are distributed as evenly as possible across these groups.
`list_categorical_column_names`required	—	A list of column names in the DataFrame to use for stratification. The function ensures the proportions of these categories are roughly equal in every resulting group.
`random_seed`	`412`	The seed value for the random number generator to ensure reproducibility.

Returns

The stratified DataFrame with an additional group column (integer values) indicating the assigned group number. Note that the DataFrame may be resorted by the stratification variables.

Example

python

from analysistoolbox.data_processing import CreateStratifiedRandomSampleGroups
import pandas as pd

# Create balanced A/B test groups based on tier and gender
users = pd.DataFrame({
    'user_id': range(1, 13),
    'tier': ['Gold', 'Gold', 'Silver', 'Silver', 'Bronze', 'Bronze'] * 2,
    'gender': ['M', 'F'] * 6
})
stratified_groups = CreateStratifiedRandomSampleGroups(
    users,
    number_of_groups=2,
    list_categorical_column_names=['tier', 'gender']
)