在 numpy 中用 3d 数组索引 2d 数组

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”，一个 3d 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”进行3次索引，以获得一个对象，如下所示：

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

您可以使用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) 的数组。这实质上是为了让它可以使用idx 进行广播，其形状为(3, 2, 5)。

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

Answer 2

可以使用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)