在pandas中为值生成唯一的零基ID

Question

我有一些数据在一个数据框里，这个数据框里有一列是标识符。

data = DataFrame({'id' : [50,50,30,10,50,50,30]})

对于每一个独特的标识符，我想生成一个新的独特标识符。我希望这些标识符是从0开始的连续整数。以下是我目前的做法：

unique = data[['id']].drop_duplicates()   
unique['group'] = np.arange(len(unique))
unique.set_index('id')
data = data.merge(unique, 'inner', on = 'id')

这个方法可以用，但感觉有点不太干净。有没有更好的办法呢？

Answer 1

这就是 pandas.factorize 的作用：

data = pd.DataFrame({'id' : [50,50,30,10,50,50,30]})
print pd.factorize(data.id)[0]

输出结果：

[0 0 1 2 0 0 1]

numpy.unique 也可以做到这一点：

import numpy as np
print np.unique([50,50,30,10,50,50,30], return_inverse=True)[1]

输出结果：

array([2, 2, 1, 0, 2, 2, 1])

numpy.unique 输出的索引是按照值排序的，所以最小的值 10 被分配给索引 0。如果你想用 factorize 得到这样的结果，可以把 sort 参数设置为 True：

pandas.factorize(data.id, sort=True)[0]