时间戳

首页 > 解决方案 > 使用 Numpy 基于向量从 3D 矩阵中选择列

python - 使用 Numpy 基于向量从 3D 矩阵中选择列

问题描述

所以我试图根据使用 Numpy 的向量中的值从 3D 矩阵中选择一些列。我已经使用列表理解解决了这个问题，但我认为使用 Numpy 的内置方法可能会有更好的方法。有谁知道这种方法或方法组合是否存在？

matrix1 = np.array([[1, 2, 3],
                    [4, 5, 6],
                    [7, 8, 9]])
matrix2 = np.array([[10, 11, 12],
                    [13, 14, 15],
                    [16, 17, 18]])
total_matrix = np.array([matrix1, matrix2])
vector = [0,1,1]

# Retrieve the first column from the first matrix, second and third from the second matrix.
result = np.array([total_matrix[index2,: , index1] for index1, index2 in enumerate(vector)]).transpose()

# result:
np.array([[1, 11, 12],
          [4, 14, 15],
          [7, 15, 18]])

标签： pythonnumpymatrix

解决方案

In [58]: total_matrix[vector, np.arange(3)[:,None], np.arange(3)]
Out[58]: 
array([[ 1, 11, 12],
       [ 4, 14, 15],
       [ 7, 17, 18]])

vector索引第一个维度。另外2个广播用它来选择需要的(3,3)。虽然我知道一般原则，但在找到正确的一种之前，我尝试了多种变体（大约 9 种）。

在另一个答案中的使用diagonal相当于做：

In [61]: total_matrix[vector][:,np.arange(3),np.arange(3)]
Out[61]: 
array([[ 1,  5,  9],
       [10, 14, 18],
       [10, 14, 18]])

推荐阅读