使用 Numpy 從基於向量的 3D 矩陣中選擇列

Question

所以我試圖根據使用 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]])

Answer 1

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

Answer 2

您可以使用您的vector對total_matrix進行切片，然后選擇它的適當對角線元素：

>>> np.diagonal(total_matrix[vector], axis1=0, axis2=2)
array([[ 1, 11, 12],
       [ 4, 14, 15],
       [ 7, 17, 18]])