如何用另一个ndarray索引一个ndarray？

Question

我正在python / numpy中做一些机器学习的东西，我想用一维ndarray索引一个二维ndarray，以便获得带有索引值的一维数组。

我用一些丑陋的代码来工作，我想知道是否有更好的方法，因为对于像python + numpy这样的好语言和模块组合来说，这似乎是不自然的。

   a = np.arange(50).reshape(10, 5) # Array to be indexed
   b = np.arange(9, -1, -2) # Indexing array
   print(a)
   print(b)
   print(a[b, np.arange(0, a.shape[1]).reshape(1,a.shape[1])])
   #Prints:
   [[ 0  1  2  3  4]
    [ 5  6  7  8  9]
    [10 11 12 13 14]
    [15 16 17 18 19]
    [20 21 22 23 24]
    [25 26 27 28 29]
    [30 31 32 33 34]
    [35 36 37 38 39]
    [40 41 42 43 44]
    [45 46 47 48 49]]

   [9 7 5 3 1]

   [[45 36 27 18  9]]

这正是我想要的（即使从技术上讲是2D ndarray），但这似乎非常复杂。 有没有更整洁的方式？

编辑：为澄清起见，我实际上不想要一维数组，这很难解释。 实际上，我确实想要一个长度为1的维度，因为这是以后处理它所需要的，但是使用reshape()语句可以轻松实现。 抱歉，我只是将实际代码与更笼统的问题混在一起。

Answer 1

您可以只使用np.diagonal来获取所需的内容。 无需调整reshape或索引。 这里最棘手的部分是确定要获取的模式，该模式基本上是a[b]矩阵的对角元素。

a = np.arange(50).reshape(10, 5) # Array to be indexed
b = np.arange(9, -1, -2) # Indexing array
print (np.diagonal(a[b]))
# [45 36 27 18  9]

正如@ user2357112在评论中提到的， np.diagonal的返回是只读的。 我认为，如果您计划将值附加/修改到此最终所需列表中，将是一个问题。 如果您只想打印它们或将它们用于其他索引，应该没问题。

根据文档

从NumPy 1.9开始，它返回原始数组的只读视图。 尝试写入结果数组将产生错误。

在将来的发行版中，它将返回一个读/写视图，并且写入返回的数组将更改您的原始数组。 返回的数组将具有与输入数组相同的类型。

如果不写此函数返回的数组，则可以忽略以上所有内容。

Answer 2

您需要一个1D数组，但是包含了一个reshape调用，其唯一目的是将数组从所需的格式转换为不需要的格式。

停止重塑arange输出。 另外，您无需显式指定0起始值：

result = a[b, np.arange(a.shape[1])]