子集在蟒蛇

Question

一個二維數組我想用一個二維數組子集另一個二維陣列（它們具有相同的長度），例如二維數組：

import numpy as np 
tmp = np.array([[0.33, 0.67], [0.67, 0.33]]) 
index = np.array([[1], [0]]) 

我想是這樣的：

In[91]: np.array([tmp[i][index[i]] for i in range(len(index))]) 

Out[91]: 
array([[ 0.67], 
     [ 0.67]]) 

它的工作，但是，有沒有更聰明/更有效的方式來做到這一點？

Answer 1

您可以使用您的index陣列的shape()和inedx本身作爲列的行的索引，然後使用一個簡單的索引，以獲得預期的項目：

>>> tmp[(np.array(index.shape[::-1])-1)[:,None], index] 
array([[ 0.67], 
     [ 0.67]])