python - 2 个数组的 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

这是一个选项:

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