python – 使用numpy数组作为另一个数组的第二个dim的索引？

[weixin_38071069]

参见英文答案 > Indexing one array by another in numpy                                    3个例如,我有两个numpy数组, A = np.array( [[0,1], [2,3], [4,5]]) B = np.array( [[1], [0], [1]], dtype='int') 我想从A的每一行中提取一个元素,并且该元素由B索引,所以我想要以下结果： C = np.array( [[1], [2], [5]]) 我试过A [：,B.ravel()],但它会播放B,而不是我想要的.也看了np.take,似乎不是我的问题的正确解决方案. 但是,我可以通过转置A来使用np.choose, np.choose(B.ravel(), A.T) 但还有其他更好的解决方案

[weixin_38077297]

你可以使用NumPy's purely integer array indexing – A[np.arange(A.shape[0]),B.ravel()] 样品运行 – In [57]: A Out[57]: array([[0, 1], [2, 3], [4, 5]]) In [58]: B Out[58]: array([[1], [0], [1]]) In [59]: A[np.arange(A.shape[0]),B.ravel()] Out[59]: array([1, 2, 5]) 请注意,如果B是一维数组或此类列索引的列表,则可以使用.ravel()跳过展平操作. 样品运行 – In [186]: A Out[186]: array([[0, 1], [2, 3], [4, 5]]) In [187]: B Out[187]: [1, 0, 1] In [188]: A[np.arange(A.shape[0]),B] Out[188]: array([1, 2, 5])