如何在Pandas中按指定列等频次抽取指定行数的数据？

需求说明 #

从名为DF的Pandas DataFrame中抽取总计N行数据，要求指定列（如"Type"）下各唯一类型的抽样频次尽可能相等。若部分类型的现有数据量无法满足基础分配数，则从其他有剩余数据的类型中随机补充抽样数量。

实现代码 #

import pandas as pd
import numpy as np

# 配置参数
target_column = "Type"  # 指定目标列
sample_total = 100      # 抽样总行数

# 1. 统计目标列各类型的现有数据量
type_data_counts = DF[target_column].value_counts()
total_types = len(type_data_counts)

# 2. 计算基础分配数与剩余待分配数量
base_sample = sample_total // total_types
remaining = sample_total % total_types

# 3. 初始化各类型抽样数量
sample_sizes = pd.Series([base_sample] * total_types, index=type_data_counts.index)

# 4. 处理数据量不足的类型，重新分配差额
for typ in type_data_counts.index:
    if sample_sizes[typ] > type_data_counts[typ]:
        # 该类型只能取全部现有数据，差额计入剩余待分配数
        diff = sample_sizes[typ] - type_data_counts[typ]
        sample_sizes[typ] = type_data_counts[typ]
        remaining += diff
        
        # 获取还有剩余数据可抽取的类型
        available_types = sample_sizes[sample_sizes < type_data_counts].index
        
        # 循环分配剩余名额，直到分完或无可用类型
        while remaining > 0 and len(available_types) > 0:
            add_count = min(remaining, len(available_types))
            selected = np.random.choice(available_types, size=add_count, replace=False)
            sample_sizes[selected] += 1
            remaining -= add_count
            available_types = sample_sizes[sample_sizes < type_data_counts].index

# 5. 分配剩余基础名额（若初始分配后还有剩余）
if remaining > 0 and len(available_types) > 0:
    selected = np.random.choice(available_types, size=remaining, replace=False)
    sample_sizes[selected] += 1

# 6. 执行抽样并合并结果
sampled_list = []
for typ, size in sample_sizes.items():
    if size > 0:
        sampled = DF[DF[target_column] == typ].sample(n=size, replace=False, random_state=42)
        sampled_list.append(sampled)

# 合并后打乱顺序，避免按类型排列
final_sample = pd.concat(sampled_list, axis=0).sample(frac=1, random_state=42).reset_index(drop=True)

代码说明 #

• 先按类型数量均分抽样名额，保证基础均衡性
• 自动处理数据量不足的类型，将差额重新分配给有富余数据的类型
• 抽样后打乱结果顺序，避免输出结果按类型分组
• 设置random_state可固定随机抽样结果，方便复现测试

内容的提问来源于stack exchange，提问作者Parsh