使用Pandas为每组生成指定最小最大值的均分ID列

给定如下示例DataFrame：

import pandas as pd

min_id = 1
max_id = 10

data = [['A', 2], ['A', 3], ['A', 1], ['A', 4], ['A', 4], ['A', 2],
        ['B', 4], ['B', 5], ['B', 7], ['B', 4], ['B', 2],
        ['C', 1], ['C', 3], ['C', 2], ['C', 1], ['C', 5], ['C', 2] ,['C', 1],
        ['D', 1], ['D', 1], ['D', 1], ['D', 1]]
df = pd.DataFrame(data = data, columns = ['group', 'val'])

需要新增一列id，要求每个分组内的id最小值为min_id（1）、最大值为max_id（10），中间值按组内行数均匀分配，最终输出如下：

data = [['A', 2, 1], ['A', 3, 2.8], ['A', 1, 4.6], ['A', 4, 6.4], ['A', 4, 8.2], ['A', 2, 10],
        ['B', 4, 1], ['B', 5, 3.25], ['B', 7, 5.5], ['B', 4, 7.75], ['B', 2, 10],
        ['C', 1, 1], ['C', 3, 2.5], ['C', 2, 4], ['C', 1, 5.5], ['C', 5, 7], ['C', 2, 8.5] ,['C', 1, 10],
        ['D', 1, 1], ['D', 1, 4], ['D', 1, 7], ['D', 1, 10]]
df_desired = pd.DataFrame(data = data, columns = ['group', 'val', 'id'])

解决方案 #

通过groupby结合组内索引的线性映射实现，代码如下：

import pandas as pd
import numpy as np

min_id = 1
max_id = 10

# 生成id列
df['id'] = df.groupby('group').apply(
    lambda x: min_id + (max_id - min_id) * x.cumcount() / (len(x)-1) if len(x) > 1 else min_id
).reset_index(level=0, drop=True)

# 格式化小数位数以匹配示例（可选）
df['id'] = df['id'].round(2)

# 查看结果
print(df)

代码说明 #

• groupby('group')：按group列分组处理每个子数据集
• cumcount()：获取每个分组内的行索引（从0开始）
• 线性映射公式：将索引值转换到[min_id, max_id]区间，公式为：

  id = min_id + (max_id - min_id) * 组内索引 / (组内行数 - 1)
  

  当组内仅1行时，直接赋值为min_id（可根据需求改为max_id）
• reset_index(level=0, drop=True)：移除分组后的多级索引，将结果合并回原DataFrame
• round(2)：格式化小数位数，与示例输出一致

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