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 长度相同的数组 b 中每个元素在 a 中的索引？我可以通过循环遍历 b 的元素 b[i] 并检查 np.where(a == b[i]) 来实现，但想知道是否有更快、更好、更短的 numpy 方法。

- rwolst

---

你想要的 "c" 是索引加1，还是仅仅是索引本身？ - mauve

啊，是的，那里有一个索引错误，我会编辑问题。应该是实际的索引。 - rwolst

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

- Akavall

---

如果a中的元素可能不在b中，则结果= np.vectorize(my_dict.get)(a，-1) - Paul

vectorize只是一个方便的函数，其底层实现是使用Python for循环。它可以节省开发人员的时间，但并不能节省CPU时间。 - Davidmh

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

- Jaime

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

这比循环并执行 np.where(a == b[i]) 更有效率，因为你只访问了数组中的一个元素。

- jabaldonedo

Answer 4

完全基于numpy的解决方案：

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

- gg349

---

有趣的解决方案！似乎也是最快的。不幸的是，它不容易阅读... - koffein

这是正确的解决方案。如果你按照PEP8标准并且没有考虑从numpy导入，那就完美了 :) - eickenberg

实际上，np.where(a[:, np.newaxis] == b)[1] 也可以完成这个任务。 - eickenberg

@eickenberg +1。我想避免使用where，因为我认为它会更慢。结果证明：对于小数组，where实际上更快；对于大数组，它稍微慢一些。 - gg349

Answer 5

另一种解决方案可能是通过以下方式：

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)

- oschoudhury

Answer 6

由于数组b包含唯一的元素，因此与a的相等性只能与b中的一个单独元素相匹配。如果所有a的元素都确定在b中，则

import numpy as np
indices = np.where(a[:, np.newaxis] == b)[1]

这将解决问题。如果您不确定a的所有元素是否都在b中，则

in_b, indices = np.where(a[:, np.newaxis] == b)

将会收集所有在b中的a元素到in_b中

- eickenberg