Numpy索引的2个数组

Question

考虑两个numpy数组

a = np.array(['john', 'bill', 'greg', 'bill', 'bill', 'greg', 'bill'])
b = np.array(['john', 'bill', 'greg'])

我怎样才能生成第三个数组

c = np.array([0,1,2,1,1,2,1])

与a表示a数组中每个条目的索引的长度相同b？

我可以通过循环遍历bas b[i]和elements的元素来看到一种方法,np.where(a == b[i])但是想知道numpy是否可以通过更快/更好/更少的代码行来实现这一点.

Answer 1

这是一个选项:

import numpy as np

a = np.array(['john', 'bill', 'greg', 'bill', 'bill', 'greg', 'bill'])
b = np.array(['john', 'bill', 'greg'])

my_dict = dict(zip(b, range(len(b))))

result = np.vectorize(my_dict.get)(a)

结果:

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

Answer 2

排序是使用numpy进行矢量化的一个很好的选择:

>>> s = np.argsort(b)
>>> s[np.searchsorted(b, a, sorter=s)]
array([0, 1, 2, 1, 1, 2, 1], dtype=int64)

如果你的数组a有m元素并且b有n,那么排序将是O(n log n),并且搜索O(m log n),这是不错的.基于字典的解决方案应该是线性分摊的,但如果数组不是很大,Python循环可能会使它们比这更慢.基于广播的解决方案具有二次复杂性,对于非常小的阵列,它们只会更快.

您的样本的一些时间安排:

In [3]: %%timeit
   ...: s = np.argsort(b)
   ...: np.take(s, np.searchsorted(b, a, sorter=s))
   ...: 
100000 loops, best of 3: 4.16 µs per loop

In [5]: %%timeit
   ...: my_dict = dict(zip(b, range(len(b))))
   ...: np.vectorize(my_dict.get)(a)
   ...: 
10000 loops, best of 3: 29.9 µs per loop

In [7]: %timeit (np.arange(b.size)*(a==b[:,newaxis]).T).sum(axis=-1)
100000 loops, best of 3: 18.5 µs per loop