如何從矩陣中選擇特定的列索引？

Question

我有一個矩陣和要從矩陣中為每一行選擇的列索引列表。 我該如何在numpy中做到這一點？

my_matrix = np.array([[1, 2], [4, 5]])
col_idx = np.array([1, 0])
selected = .... # selects 1st element of row 0 and 0th element of row 1.
print selected
# np.array([2, 4])

Answer 1

您可以使用范圍切片：

In [11]: my_matrix[np.arange(my_matrix.shape[0]), col_idx]
Out[11]: array([2, 4])

Answer 2

np.choose對於進行以下選擇非常有用：

>>> np.choose(col_idx, my_matrix.T)
array([2, 4])

在更大的矩陣上：

>>> my_matrix_2 = np.array([[1, 2], [4, 5], [3, 7], [4, 1]])
>>> col_idx_2 = np.array([1, 0, 0, 1])
>>> np.choose(col_idx_2, my_matrix_2.T)
array([2, 4, 3, 1])

該方法返回具有選定值的新數組（而不是原始數組的視圖）。

---

文檔中還有更多這種（最初是不太明顯的）方法的示例，但是我將使用上面的第二個示例來解釋發生了什么。

我們正在使用np.choose從名為my_matrix_2.T的選擇數組中返回一個新數組，其中col_idx_2指定我們每次應從選擇數組的哪一行中選擇。

注意，我們轉置了my_matrix_2使其工作：

# my_matrix_2.T
array([[1, 4, 3, 4],  # row 0
       [2, 5, 7, 1]]) # row 1

我們有col_idx_2 = [1, 0, 0, 1] 。 現在一次遍歷此數組一個值：

• 新陣列的第一個元素將是行的第一元件1的my_matrix_2.T 。 這是2 。

• 新陣列的第二個元素將是行的第二元件0的my_matrix_2.T 。 這是4 。

• 新陣列的第三個元素將是行的第三元件0的my_matrix_2.T 。 這是3 。

• 新的數組的第四個元素將是行的第四個元素1的my_matrix_2.T 。 這是1 。

因此，該方法返回array([2, 4, 3, 1]) 。

Answer 3

In [211]: M = np.array([[1, 2], [4, 5]])

In [212]: cid = [1, 0]

In [213]: M[[list(i) for i in zip(range(M.shape[0]), cid)]]
Out[213]: array([2, 4])