用3D数组索引2D数组的numpy方法

Question

我有两个数组。

"a"是一个二维的numpy数组。

import numpy.random as npr

a = array([[5,6,7,8,9],[10,11,12,14,15]])
array([[ 5,  6,  7,  8,  9],
       [10, 11, 12, 14, 15]])

"idx"是一个三维的numpy数组，里面包含了我想用来索引"a"的三种索引变体。

idx = npr.randint(5, size=(nsamp,shape(a)[0], shape(a)[1]))
array([[[1, 2, 1, 3, 4],
        [2, 0, 2, 0, 1]],

       [[0, 0, 3, 2, 0],
        [1, 3, 2, 0, 3]],

       [[2, 1, 0, 1, 4],
        [1, 1, 0, 1, 0]]])

现在我想用"idx"中的索引三次来获取"a"中的数据，结果应该是这样的：

array([[[6, 7, 6, 8, 9],
        [12, 10, 12, 10, 11]],

       [[5, 5, 8, 7, 5],
        [11, 14, 12, 10, 14]],

       [[7, 6, 5, 6, 9],
        [11, 11, 10, 11, 10]]])

直接用"a[idx]"是行不通的。有没有什么办法可以做到这一点？（我使用的是Python 3.4和numpy 1.9）

Answer 1

你可以使用 take 这个数组方法：

import numpy

a = numpy.array([[5,6,7,8,9],[10,11,12,14,15]])

idx = numpy.random.randint(5, size=(3, a.shape[0], a.shape[1]))

print a.take(idx)

Answer 2

你可以使用 choose 来从 a 中进行选择：

>>> np.choose(idx, a.T[:,:,np.newaxis])
array([[[ 6,  7,  6,  8,  9],
        [12, 10, 12, 10, 11]],

       [[ 5,  5,  8,  7,  5],
        [11, 14, 12, 10, 14]],

       [[ 7,  6,  5,  6,  9],
        [11, 11, 10, 11, 10]]])

如你所见，a 需要先从一个形状为 (2, 5) 的数组调整为形状为 (5, 2, 1) 的数组。这么做的主要原因是为了让它能够和形状为 (3, 2, 5) 的 idx 进行广播（也就是可以一起运算）。

（我从 @immerrr 的回答中学到了这个方法，链接在这里： https://stackoverflow.com/a/26225395/3923281）