使用numpy中的3d数组索引2d数组

Question

我有两个数组。

“a”，一个2d 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”索引三次以获得一个对象，如下所示：

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) 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)