首页 > 解决方案 > 给定每一行对应的索引,从一个矩阵中获取对应的元素

问题描述

给定每行的索引,如何返回二维矩阵中的相应元素?例如,在数组中,np.array([[1,2,3,4],[4,5,6,7]])我希望看到[[1,2],[4,5]]给定的输出indxs = np.array([[0,1],[0,1]])。以下是我尝试过的:

a= np.array([[1,2,3,4],[4,5,6,7]])
indxs = np.array([[0,1],[0,1]]) #means return the elements located at 0 and 1 for each row
#I tried this, but it returns an array with shape (2, 2, 4) 
a[idxs]

标签: pythonarraysnumpy

解决方案


您获得两倍数组的原因是,当您这样做时,a[[0,1]]您正在从数组 a 中选择第 0 行和第 1 行,这确实是您的整个数组。

In[]: a[[0,1]]

Out[]: array([[1, 2, 3, 4],
       [4, 5, 6, 7]])

您可以使用幻灯片获得所需的输出。那将是最简单的方法。

a = np.array([[1,2,3,4],[4,5,6,7]])
a[:,0:2]

Out []: array([[1, 2],
               [4, 5]])

如果您仍然对索引感兴趣,您还可以让您的输出:

In[]: [list(a[[0],[0,1]]),list(a[[1],[0,1]])]

Out[]: [[1, 2], [4, 5]]

NumPy文档为您提供了关于索引如何工作的非常好的概述。


推荐阅读