Python Pandas DataFrame新增Category列报错及高效实现问询

1. 索引越界错误的解决方法 #

错误原因 #

原代码核心逻辑存在错误：

• pd.factorize(sorted(x.unique()))[0] 生成的是当前分组下排序后唯一颜色的标签数组，长度等于该分组的颜色唯一值数量（比如Citrus分组只有3个唯一颜色，数组长度为3）。
• x.argsort() 生成的是当前分组所有Color值排序后的位置索引，长度等于该分组的行数（比如Citrus分组有5行，索引范围是0-4）。
  用后者索引前者时，当分组行数大于唯一值数量，就会触发索引越界（比如索引3超出了长度为3的数组范围）。

另外注意代码拼写错误：import panda as pd 应改为 import pandas as pd。

修正后的代码 #

import pandas as pd

data = {'Fruit': ['Citrus', 'Apples', 'Apples', 'Citrus', 'Citrus', 'Apples', 'Citrus', 'Apples', 'Citrus', 'Apples'],
    'Color': ['Orange', 'Dark Red', 'Red', 'Orange', 'Green', 'Red', 'Yellow', 'Dark Red', 'Yellow', 'Green']}

df = pd.DataFrame(data)

# 定义分组内颜色到分类索引的映射逻辑
def map_color_to_category(x):
    # 获取当前分组排序后的唯一颜色列表
    sorted_unique_colors = sorted(x.unique())
    # 创建颜色到索引的字典
    color_index_map = {color: idx for idx, color in enumerate(sorted_unique_colors)}
    # 对当前分组的Color列进行映射
    return x.map(color_index_map)

# 按Fruit分组后应用映射逻辑生成Category列
df['Category'] = df.groupby('Fruit')['Color'].transform(map_color_to_category)

print(df)

运行后输出符合预期：

Fruit     Color  Category
0    Citrus    Orange         1
1    Apples  Dark Red         0
2    Apples       Red         2
3    Citrus    Orange         1
4    Citrus     Green         0
5    Apples       Red         2
6    Citrus    Yellow         2
7    Apples  Dark Red         0
8    Citrus    Yellow         2
9    Apples     Green         1

2. 10万行数据的高效实现方式 #

针对大数据量，推荐以下两种更高效的方案：

方案一：利用pd.Categorical生成编码 #

pd.Categorical的codes属性可直接生成分类索引，结合分组transform操作，性能优异：

df['Category'] = df.groupby('Fruit')['Color'].transform(
    lambda x: pd.Categorical(x, categories=sorted(x.unique())).codes
)

方案二：先构建映射表再合并 #

预先生成所有(Fruit, Color)对应的Category索引，再与原表合并，避免逐组循环，适合超大数据量：

# 生成(Fruit, Color)到Category的映射表
color_category_map = df.groupby('Fruit')['Color'].unique()\
                      .apply(sorted)\
                      .explode()\
                      .reset_index()
color_category_map['Category'] = color_category_map.groupby('Fruit').cumcount()

# 合并映射表到原DataFrame
df = df.merge(color_category_map, on=['Fruit', 'Color'], how='left')

这两种方案时间复杂度更低，在10万行数据上的运行效率远高于自定义map函数的方式。

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