Scenario

A customer dataset contains records from multiple regions with varying sizes. You need to create a representative sample that maintains the proportional distribution of each region.

Task

Write a Python script at /home/interview/stratified_sample.py using pandas that reads /home/interview/customers.csv, extracts a 10% stratified sample (proportionally representative from each region), and writes the result to /home/interview/sample_data.csv.

Note: pandas is already installed.

Step 1: Examine the data distribution

python3 -c "import pandas as pd; df = pd.read_csv('/home/interview/customers.csv'); print(df['region'].value_counts().sort_index())"

Shows the distribution of customers across regions.

Step 2: Create the Python script

nano /home/interview/stratified_sample.py

Write a script to perform stratified sampling:

import pandas as pd

# Read the CSV file
df = pd.read_csv('/home/interview/customers.csv')

# Perform stratified sampling - 10% from each region
# Use sample with frac parameter on grouped object
sample_df = df.groupby('region', group_keys=False).sample(frac=0.1, random_state=42)

# Reset index for clean output
sample_df = sample_df.reset_index(drop=True)

# Write to CSV
sample_df.to_csv('/home/interview/sample_data.csv', index=False)

print(f"Sampled {len(sample_df)} records (10% of {len(df)})")

Step 3: Run the script

python3 /home/interview/stratified_sample.py

Step 4: Verify the output

python3 -c "import pandas as pd; df = pd.read_csv('/home/interview/sample_data.csv'); print(f'Total: {len(df)} records'); print(df['region'].value_counts().sort_index())"

Should show 200 total records with proportional representation from each region.