NumPy数组花式索引

Question

我有一个（9099，2，26）形状的numpy数组。

我想从这个数组中得到一个形状为（9099，2，1）的新数组，其中第3维的值将由另一个1d数组（9099，）定义。 这个1d数组的每个值实际上是我的主数组中第3维所需值的索引。

例如，主数组的第9个元素如下所示：

In [1]: a[9,:,:] Out [1]: array([[ 0, 1011165, 1098400, 1140065, 1172240, 1203971, 1439189, 1600540, 2281575, 8888888, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [ 280000, 263998, 974020, 653653, 888965, 851341, 466195, 424801, 0, 9999999, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])

1d数组的第9个元素是：

In [2]: b[9] Out [2]: 9

因此，所需输出的第十个元素应如下所示：

In [3]: c[9,:,:] Out [3]: array([[8888888], [9999999]])

我想应该通过花哨的索引来完成它？

谢谢！

Answer 1

您可以使用

a[np.arange(len(index)), :, index]