Pandas如何根据字典指定行数为DataFrame新增分组列

直接通过等长序列赋值即可实现需求，不需要复杂分箱或容易错位的循环，核心是生成和DataFrame行数完全一致、每个键重复次数匹配字典值的序列，直接赋值给新列即可。 #

首先对齐测试数据：

import pandas as pd
import numpy as np

# 初始DataFrame
df = pd.DataFrame({
    'ID': [1,2,3,4,5,7,6,8,9,10],
    'val': ['a','b','c','d','a','d','v','j','k','a']
})
# 分块配置字典
split_config = {'aa':3, 'bb':3, 'cc':4}

最简实现（推荐） #

用numpy.repeat直接生成目标列，一行代码完成：

df['new_col'] = np.repeat(
    a=list(split_config.keys()),
    repeats=list(split_config.values())
)

运行后输出结果完全匹配需求：

ID val new_col
0   1   a      aa
1   2   b      aa
2   3   c      aa
3   4   d      bb
4   5   a      bb
5   7   d      bb
6   6   v      cc
7   8   j      cc
8   9   k      cc
9  10   a      cc

如果需要随机分配类别、不按原行顺序从上到下匹配，只需要先打乱DataFrame顺序再执行赋值即可：

# 随机打乱行顺序，random_state可固定随机结果
df = df.sample(frac=1, random_state=42).reset_index(drop=True)
df['new_col'] = np.repeat(list(split_config.keys()), list(split_config.values()))

之前方案踩坑原因说明 #

• 遍历字典取数错位：本质是没有维护行偏移量，每次取数没有跳过已经分配过类别的行，如果坚持用循环写法，可以参考以下无错位实现：

new_col_values = []
for key, count in split_config.items():
    new_col_values.extend([key]*count)
df['new_col'] = new_col_values

• pd.cut报单调递增错误：pd.cut要求传入的bins是从最小值到最大值、单调递增的分箱边界列表，直接传入字典的count值（本例为[3,3,4]）既没有从0开始、存在重复值，也没有计算累计边界，自然会触发报错。如果要用pd.cut实现，需要先计算累计分箱边界：

# 生成分箱边界 [0,3,6,10]
bins = [0] + list(np.cumsum(list(split_config.values())))
df['new_col'] = pd.cut(
    x=df.index,
    bins=bins,
    labels=list(split_config.keys())
)

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