使用一个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..希望这是有意义的！