用一批 numpy 索引矩阵索引一批 numpy 向量

Question

如果我有一个向量

vec

，我可以用矩阵对其进行索引，如下所示：

import numpy as np

vec = np.asarray([1,2,3,4]) # Shape (4,)

mat = np.asarray([[0,2],
                  [3,1]]) # Shape (2,2)

result = vec[mat] # Shape (2,2)


# result = array([[1, 3],
#              [4, 2]])

现在假设我有一批向量和矩阵。 严格使用 numpy 对批次中的每个元素进行相同类型的索引的最简单方法是什么？例如：

vecs = np.asarray([[1, 2, 3, 4],
                   [5, 6, 7, 8],
                   [9,10,11,12]]) # Shape (3,4)

mats = np.asarray([ [[0,1],
                     [1,0]],
                    [[0,2],
                     [1,1]],
                    [[3,1],
                     [0,0]] ]) # Shape (3,2,2)

# results = np.asarray([ [[1,2],
#                         [2,1]],
#                        [[5,7],
#                         [6,6]],
#                        [[12,10],
#                         [9,9]] ]) # Shape (3,2,2)

Answer 1

添加 (3,1,1) 数组来索引第一个维度：

In [196]: vecs[np.arange(3)[:,None,None],mats]
Out[196]: 
array([[[ 1,  2],
        [ 2,  1]],

       [[ 5,  7],
        [ 6,  6]],

       [[12, 10],
        [ 9,  9]]])

这将与您的 (3,2,2)

mats

一起广播。