NumPy中的快速向量化索引

Question

假设我们有另一个numpy数组的索引数组：

import numpy as np
a = np.array([0, 3, 1])
b = np.array([0, 10, 20, 30, 40, 50, 60, 70])

我们可以直接使用数组 a 作为索引:

b[a] # np.array([0, 30, 10])

但是如果数组a有多个维度怎么办？例如，

a = np.array([[0, 2], [1, 3], [2, 4]])
# I want to get b[a] = np.array([[0, 20], [10, 30], [20, 40]])

如果数组a的维度大于1，则Numpy索引无法工作。我们可以通过使用map来实现所需的结果。

map(lambda x: b[x], a)

然而，它的速度相当慢。对于一维情况，直接索引比使用map快10-100倍。

有没有更快的方法？

- DLunin

---

b[a] 会发生什么？ - hpaulj

即使 a 是2D，b[a]也会按预期工作。你遇到了什么错误？ - jakevdp

Answer 1

问题是什么？我可以用2D数组索引b。输出的形状与a1匹配：

In [64]: b = np.array([0, 10, 20, 30, 40, 50, 60, 70])
In [65]: a1 = np.array([[0, 2], [1, 3], [2, 4]])
In [66]: b[a1]
Out[66]: 
array([[ 0, 20],
       [10, 30],
       [20, 40]])

b[a1]并不等同于b[a1[:,0],a1[:,1]]。这是因为a1的两列并不能提供两个索引（这需要一个2D的b数组）。

- hpaulj

---

再次检查了一下。你是对的：如果a1是numpy数组，它确实可以工作。然而，如果a1是列表（[[0, 2], [1, 3], [2, 4]]），它就不行，这也是为什么一开始会有这个问题的原因。我现在会修正这个问题。 - DLunin

你的答案就在这里：b[np.array(a)]。b.take(a)稍微快一点。 - hpaulj

Answer 2

有内置的np.take可用于完成同样的任务 -

np.take(b,a)

你可以使用 .ravel() 将 a 扁平化，然后索引到 b 并重新调整回到 a 的形状 -

b[a.ravel()].reshape(a.shape)

这些基于NumPy的方法在性能和内存方面都比使用map(lambda x: b[x], a)更好，因为使用map会给我们提供数组列表。

示例运行 -

In [34]: a
Out[34]: 
array([[0, 2],
       [1, 3],
       [2, 4]])

In [35]: b
Out[35]: array([ 0, 10, 20, 30, 40, 50, 60, 70])

In [36]: np.take(b,a)
Out[36]: 
array([[ 0, 20],
       [10, 30],
       [20, 40]])

In [37]: b[a.ravel()].reshape(a.shape)
Out[37]: 
array([[ 0, 20],
       [10, 30],
       [20, 40]])

运行时测试 -

In [39]: a = np.random.randint(0,100,(200,100))

In [40]: b = np.random.randint(0,100,(20000))

In [41]: %timeit map(lambda x: b[x], a)
1000 loops, best of 3: 643 µs per loop

In [42]: %timeit np.take(b,a)
10000 loops, best of 3: 105 µs per loop

In [43]: %timeit b[a.ravel()].reshape(a.shape)
1000 loops, best of 3: 231 µs per loop

- Divakar

---

在这种情况下，take 做的和 b[a] 是一样的。 - hpaulj

@hpaulj 没错！你说得对，那也可以！我想如果a是一个列表时，可以使用.take()。 - Divakar

np.take 就是我一直在寻找的！ - Jaakko