Numpy：从数组中获取索引在另一个数组中的值

Question

我有一个 mx1 的数组 a，里面存了一些值。此外，我还有一个 nxk 的数组 b，里面存的是从 0 到 m 的索引。

举个例子：

a = np.array((0.1, 0.2, 0.3))
b = np.random.randint(0, 3, (4, 4))

我想根据 b 中的每个索引值，获取 a 中对应的值。我可以用循环来实现：

c = np.zeros_like(b).astype('float')
n, k = b.shape
for i in range(n):
    for j in range(k):
        c[i, j] = a[b[i, j]]

有没有什么更优雅的 numpy 函数或者技巧可以做到这一点？我觉得这种方法看起来有点笨。

顺便说一下，a 和 b 最初是 Pandas 对象，如果这有帮助的话。

Answer 1

试试用一个 numpy.flatiter 对象来进行迭代：

a = np.array((0.1, 0.2, 0.3))
b = np.random.randint(0, 3, (4, 4))

c = np.array([a[i] for i in b.flat]).reshape(b.shape)
print(c)

array([[ 0.2,  0.2,  0.2,  0.1],
       [ 0.3,  0.3,  0.2,  0.1],
       [ 0.2,  0.1,  0.3,  0.3],
       [ 0.3,  0.3,  0.3,  0.1]])

Answer 2

>>> a
array([ 0.1,  0.2,  0.3])
>>> b
array([[0, 0, 1, 1],
       [0, 0, 1, 1],
       [0, 1, 1, 0],
       [0, 1, 0, 1]])
>>> a[b]
array([[ 0.1,  0.1,  0.2,  0.2],
       [ 0.1,  0.1,  0.2,  0.2],
       [ 0.1,  0.2,  0.2,  0.1],
       [ 0.1,  0.2,  0.1,  0.2]])

好了！这就是 a[b] 的用法。（另外，你可能想把 randint 函数中的上限设置为 3。）