python - 在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”进行三次索引以获得一个对象，如下所示:

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)