快速矢量索引

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 = np.array([[0, 2], [1, 3], [2, 4]]) 
# I want to get b[a] = np.array([[0, 20], [10, 30], [20, 40]]) 

numpy的索引不起作用，如果一个维数大于1，我们可以通过使用map

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

但是达到预期的效果，这是相当缓慢。对于一维情况，直接索引比使用map快约10-100倍。

有没有办法更快地做到这一点？

Answer 1

还有就是内置np.take正是出于这个相同的任务 -

np.take(b,a) 

您也可以拼合a与.ravel()，索引b，重塑回a's形状 -

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

这些NumPy的如果使用map会给我们一个数组列表，那么在性能和内存方面都会比map(lambda x: b[x], a)好得多。

采样运行 -

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

什么问题？我可以用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）。