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

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

Answer 2

内置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

numpy中的快速矢量化索引

2 个答案: