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

Question

使用一个numpy数组作为另一个数组的第二个DUMP的索引？例如，我有两个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似乎不是解决我问题的正确办法。但是，我可以用np.choose转置A,np.choose(B.ravel(), A.T)但还有其他更好的解决办法吗？

Answer 1

你可以用NumPy's purely integer array indexing -

A[np.arange(A.shape[0]),B.ravel()]

样本运行-

In [57]: AOut[57]: array([[0, 1],
       [2, 3],
       [4, 5]])In [58]: BOut[58]: array([[1],
       [0],
       [1]])In [59]: A[np.arange(A.shape[0]),B.ravel()]Out[59]: array([1, 2, 5])

请注意如果B是1D数组或此类列索引的列表，您可以简单地跳过以下操作.ravel().

样本运行-

In [186]: AOut[186]: array([[0, 1],
       [2, 3],
       [4, 5]])In [187]: BOut[187]: [1, 0, 1]In [188]: A[np.arange(A.shape[0]),B]Out[188]: array([1, 2, 5])

Answer 2

C = np.array([A[i][j] for i,j in enumerate(B)])

Answer 3

如果我从问题中得到了预期的输出，你就不需要在这里广播了。我们基本上是使用整数在每个模糊，选择元素。我们有第二个模糊指数B，所以我们只需要为第一个DIM创建相应的np.arange..希望这是有意义的！