如何在Pandas中对分组列进行可扩展的重新编号？

需求概述 #

需要一种可扩展的方法，对Pandas DataFrame中的cluster列进行重新编号，将原有的离散聚类值替换为从1开始的连续整数编号。数据规模可达数百万行、数百个聚类，手动处理不可行，此前尝试通过groupby("cluster")迭代分组分配编号未成功。

原数据结构 #

import pandas as pd

data = [
    ["Adam", 1, 22],
    ["Eddy", 0, 22],
    ["Boby", 9, 22],
    ["Timy", 2, 26],
    ["Carl", 9, 26],
    ["Anna", 0, 33],
    ["Paul", 5, 35],
    ["Mike", 7, 51],
]

current = pd.DataFrame(data,
                       columns=['name', 'score', 'cluster']
                       ).astype({'name':'string', 'score':'Int64', 'cluster':'Int64'})

当前表格状态

+----+--------+---------+-----------+
|    | name   |   score |   cluster |
|----+--------+---------+-----------|
|  0 | Adam   |       1 |        22 |
|  1 | Eddy   |       0 |        22 |
|  2 | Boby   |       9 |        22 |
|  3 | Timy   |       2 |        26 |
|  4 | Carl   |       9 |        26 |
|  5 | Anna   |       0 |        33 |
|  6 | Paul   |       5 |        35 |
|  7 | Mike   |       7 |        51 |
+----+--------+---------+-----------+

目标表格状态

+----+--------+---------+-----------+
|    | name   |   score |   cluster |
|----+--------+---------+-----------|
|  0 | Adam   |       1 |         1 |
|  1 | Eddy   |       0 |         1 |
|  2 | Boby   |       9 |         1 |
|  3 | Timy   |       2 |         2 |
|  4 | Carl   |       9 |         2 |
|  5 | Anna   |       0 |         3 |
|  6 | Paul   |       5 |         4 |
|  7 | Mike   |       7 |         5 |
+----+--------+---------+-----------+

高效可扩展解决方案 #

以下两种方法均为矢量化操作，处理百万级数据效率极高：

方法1：使用pd.factorize()

factorize()会自动将唯一的聚类值映射为从0开始的连续整数，加1即可得到从1开始的编号：

current['cluster'] = pd.factorize(current['cluster'])[0] + 1

方法2：转换为分类类型后使用cat.codes

先将cluster列转为分类类型，再通过cat.codes获取连续编号，加1调整起始值：

current['cluster'] = current['cluster'].astype('category').cat.codes + 1

原理说明 #

这两种方法都避免了循环迭代，直接利用Pandas的矢量化运算能力，处理大规模数据时性能远优于groupby迭代的方式。它们会自动识别cluster列中的唯一值，并按首次出现的顺序分配连续编号，完全符合需求。

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