Numpy 两个数组的索引

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 中的索引呢？

我想到一种方法，就是遍历 b 的每个元素，像这样 b[i]，然后用 np.where(a == b[i]) 来检查，但我在想，numpy 有没有更快、更好或者更简洁的方式来做到这一点呢？

Answer 1

还有另一种解决方案可以尝试：

arr, bSorted, ind =  np.unique(a, return_index=True, return_inverse=True)
c = bSorted[ind]

如果你想从 a 中提取出唯一的元素，并且不在乎 b 中的顺序，也就是说 b 和 c 的样子可能会不同，那么可以简化为：

b, c = np.unique(a, return_inverse=True)

Answer 2

完全使用numpy的解决方案：

(arange(b.size)*(a==b[:,newaxis]).T).sum(axis=-1)

Answer 3

首先，创建一个字典，用来把每个字符串转换成对应的数字。然后，使用 numpy.vectorize 来生成你想要的输出数组。

>>> import numpy as np
>>> a = np.array(['john', 'bill', 'greg', 'bill', 'bill', 'greg', 'bill'])
>>> b = np.array(['john', 'bill', 'greg'])
>>> d = {k:v for v, k in enumerate(b)}
>>> c = np.vectorize(d.get)(a)
>>> c
 array([0, 1, 2, 1, 1, 2, 1])

这样做比一个一个循环检查要高效得多，因为你只需要访问数组中的每个元素一次。

Answer 4

排序是使用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

Answer 5

这里有一个选择：

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])