numpy使用多維索引數組在另一個多維數組

Question

我有一個2多維數組，我想用一個作爲索引來產生一個新的多維數組。例如：

a = array([[4, 3, 2, 5], 
      [7, 8, 6, 8], 
      [3, 1, 5, 6]]) 

b = array([[0,2],[1,1],[3,1]]) 

我要使用的第一陣列中b鍵的第一陣列中返回這些索引元素，依此類推。所以我想輸出爲：

array([[4,2],[8,8],[6,1]]) 

這可能很簡單，但我無法通過搜索找到答案。謝謝。

Answer 1

這是一個有點棘手，但下面將做到這一點：

>>> a[np.arange(3)[:, np.newaxis], b] 
array([[4, 2], 
     [8, 8], 
     [6, 1]]) 

你需要指數雙雙行和a陣列的列，所以以匹配您b數組，你需要一個數組像這樣：

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

然後a[rows, b]會清楚地返回你以後的內容。您可以依靠上述廣播獲得相同的結果，用np.arange(3)[:, np.newaxis]代替rows陣列，相當於np.arange(3).reshape(3, 1)。