子集在蟒蛇

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