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 倍。

有没有更快的方法？

Answer 1

完全相同的任务有内置 np.take -

np.take(b,a)

您还可以用 .ravel() 压平 a，索引到 b 并重塑回 a's 形状 -

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

Answer 2

有什么问题吗？我可以用二维数组索引 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 的 2 列不提供两个索引（这需要 2d b）。

numpy 中的快速矢量化索引

Fast vectorized indexing in numpy

arrays

performance

numpy

vectorization

python-2.x