如何用Python按列B的8:1:1比例拆分数据集为Train/Test/Val

核心逻辑 #

针对每行的Column B计数，按8:1:1比例分配Train/Test/Val的数量，余数按顺序依次分配给三个划分（优先Train，再Test，最后Val），保证总数和原Column B完全一致。之后生成每条数据的明细行，标注对应的数据集划分。

代码实现 #

先导入依赖并构造示例数据：

import pandas as pd

# 替换成你的实际数据
df = pd.DataFrame({
    'Column A': ['类别1', '类别2'],
    'Column B': [23, 20]
})

定义比例拆分函数，处理余数分配：

def split_by_ratio(total):
    # 定义8:1:1的比例
    ratios = [8, 1, 1]
    total_ratio = sum(ratios)
    # 计算基础分配数量
    base_num = total // total_ratio
    # 计算剩余未分配的数量
    remainder = total % total_ratio
    
    counts = [base_num * r for r in ratios]
    # 把余数按顺序分配到各个分组
    for i in range(remainder):
        counts[i % len(ratios)] += 1
    
    return {'Train': counts[0], 'Test': counts[1], 'Val': counts[2]}

生成最终的标注表格：

# 给每行计算各划分的数量
df[['Train', 'Test', 'Val']] = df['Column B'].apply(split_by_ratio).apply(pd.Series)

# 生成明细行
result_rows = []
for _, row in df.iterrows():
    # 批量添加Train行
    result_rows += [{'Column A': row['Column A'], 'Column B': 1, 'Column C': 'Train'}] * row['Train']
    # 批量添加Test行
    result_rows += [{'Column A': row['Column A'], 'Column B': 1, 'Column C': 'Test'}] * row['Test']
    # 批量添加Val行
    result_rows += [{'Column A': row['Column A'], 'Column B': 1, 'Column C': 'Val'}] * row['Val']

# 转换成DataFrame
final_df = pd.DataFrame(result_rows)

验证结果 #

运行以下代码查看分组统计，确认拆分比例符合要求：

print(final_df.groupby(['Column A', 'Column C']).agg({'Column B': 'count'}))

输出结果：

Column B
Column A Column C         
类别1     Test           2
         Train         18
         Val            3
类别2     Test           2
         Train         16
         Val            2

优化方案（大数据量场景） #

如果你的数据集行数很多，iterrows效率较低，可以用向量化操作替代：

df['split_details'] = df['Column B'].apply(split_by_ratio)
# 拆分字典为多行
df = df.explode('split_details').reset_index(drop=True)
# 提取划分类型和数量
df['Column C'] = df['split_details'].apply(lambda x: next(iter(x.keys())))
df['count'] = df['split_details'].apply(lambda x: next(iter(x.values())))
# 按数量重复行
final_df = df.loc[df.index.repeat(df['count'])].drop(['split_details', 'count'], axis=1)
final_df['Column B'] = 1

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