如何在Pandas中根据id列的分组长度创建新分类列cat？

先把你的原始数据和需求理清楚：
原始DataFrame是这样的：

import pandas as pd

df = pd.DataFrame({
    'id': [1,2,2,3,1,4,3,2,3],
    'val': [21,35,45,55,10,90,45,78,23]
})

需要新增一列cat，规则是：

• 若某个id的出现次数≤1，cat取值'A'
• 若出现次数<3（也就是次数为2），cat取值'B'
• 若出现次数≥3，cat取值'C'

你尝试写函数的思路是对的，但可能没写完，这里给你两个高效简洁的实现方法：

方法一：用groupby.transform+np.select（推荐） #

这种方法不需要自己写复杂的分组函数，Pandas的transform能直接把每个id的出现次数映射到每一行，再用np.select清晰定义条件映射：

import numpy as np

# 计算每个id的出现次数，广播到原DataFrame的每一行
id_counts = df.groupby('id')['id'].transform('size')

# 定义条件和对应分类值
conditions = [
    id_counts <= 1,
    id_counts < 3,  # 第一个条件已经过滤了≤1的情况，这里实际匹配次数为2的id
    id_counts >= 3
]
category_values = ['A', 'B', 'C']

# 生成cat列
df['cat'] = np.select(conditions, category_values)

方法二：用value_counts生成映射字典+apply #

如果更喜欢用字典映射的方式，也可以先统计每个id的出现次数，再用apply逐行判断：

# 生成id到出现次数的映射字典
count_map = df['id'].value_counts().to_dict()

# 逐行根据id对应的次数赋值
df['cat'] = df['id'].apply(lambda x: 'A' if count_map[x] <=1 else 'B' if count_map[x]<3 else 'C')

运行任意一种方法后，都能得到你想要的预期输出：

id  val cat
0   1   21   B
1   2   35   C
2   2   45   C
3   3   55   C
4   1   10   B
5   4   90   A
6   3   45   C
7   2   78   C
8   3   23   C

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