如何在不循环的情况下将唯一组映射到id

Question

假设以下numpy数组

>>> z = np.zeros(3, dtype={'names': ("id", "dim1", "cnt1"), 'formats': ('i8', 'S3', 'u8')})
>>> z
array([(0, '', 0L), (0, '', 0L), (0, '', 0L)],
dtype=[('id', '<i8'), ('dim1', 'S3'), ('cnt1', '<u8')])
>>> z["dim1"] = ["foo", "foo", "bar"]
>>>
>>>
>>> z["cnt1"] = [1,2,3]
>>> z
array([(0, 'foo', 1L), (0, 'foo', 2L), (0, 'bar', 3L)],
dtype=[('id', '<i8'), ('dim1', 'S3'), ('cnt1', '<u8')])

我想将每个唯一的dim1映射到一个id。使用唯一dim1值上的for循环的一种方法如下：

>>> unique_groups = np.unique(z["dim1"])
>>> groups = z["dim1"]
>>> for idx, ug in enumerate(unique_groups):
...     z["id"][ug == groups] = idx
...
>>> z
array([(1, 'foo', 1L), (1, 'foo', 2L), (0, 'bar', 3L)],
dtype=[('id', '<i8'), ('dim1', 'S3'), ('cnt1', '<u8')])
>>>

我想知道是否有办法在没有for循环的情况下，用矢量运算来实现它。我试着用矢量化功能来执行它，如下所示

>>> def map_column(key):
...     return m[key]
...
>>> m
{'foo': 1, 'bar': 0}
>>> f = np.vectorize(map_column, otypes=[str])
>>> f(z["dim1"])
array(['1', '1', '0'],
dtype='|S1')

还有其他更有效的方法吗？在这两种方式中，哪一种被认为性能更好？

Answer 1

您可以使用.searchsorted():

In [2]: unique_groups = np.unique(z["dim1"])
In [3]: z["id"] = unique_groups.searchsorted(z["dim1"])
In [4]: z
Out[4]:
array([(1, b'foo', 1), (1, b'foo', 2), (0, b'bar', 3)],
dtype=[('id', '<i8'), ('dim1', 'S3'), ('cnt1', '<u8')])

不确定性能，但可能不会好到哪里去。