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

（我从@ 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)

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

2 个答案: