花式索引一個numpy的矩陣：每

Question

行

一個元件I具有2D陣列numpy的，矩陣形狀（M，N）。我的實際使用情況下的具有m〜1E5和n〜100，但對於具有簡單最小示例的目的：

matrix = np.arange(5*3).reshape((5, 3)) 

我有一個整數索引陣列，IDX，形狀（M）的，與每個條目在[0，n）之間。此數組指定應從矩陣的每一行中選擇哪一列。

idx = np.array([2, 0, 2, 1, 1]) 

所以，我從第2行，第1列從第4行試圖從0行，列0選擇第2列從第1行，第2列從第1行和列1。因此最終的答案應該是：

correct_result = np.array((2, 3, 8, 10, 13)) 

我曾嘗試以下，這是直觀的，但不正確的：

incorrect_result = matrix[:, idx] 

什麼上面的語法確實是適用IDX作爲一個奇特的索引數組，逐行，導致另一形狀矩陣（m，n），這不是我想要的。

這種花式索引的正確語法是什麼？

Answer 1

correct_result = matrix[np.arange(m), idx] 

的高級索引表達matrix[I, J]給出了一個輸出，使得output[n] == matrix[I[n], J[n]]。

如果我們想output[n] == matrix[n, idx[n]]，那麼我們需要I[n] == n和J[n] == idx[n]，所以我們需要I是np.arange(m)和J是idx。