python - 将 numpy 索引应用于矩阵

Question

我花了最后一个小时试图解决这个问题

假设我们有

import numpy as np
a = np.random.rand(5, 20) - 0.5
amin_index = np.argmin(np.abs(a), axis=1)
print(amin_index)
> [ 0 12  5 18  1] # or something similar

这不起作用:

a[amin_index]

因此，本质上，我需要找到数组 np.abs(a) 沿某个轴的最小值，然后从数组 a 的这些位置提取值。如何将索引仅应用于一个轴？

可能很简单，但我无法弄清楚。另外，我不能使用任何循环，因为我必须对具有数百万个条目的数组执行此操作。谢谢😊

Answer 1

一种方法是传入行索引数组(例如 [0,1,2,3,4])和每个相应行中最小值的列索引列表(您的列出amin_index)。

这将返回一个数组，其中包含每行i的[i, amin_index[i]]处的值:

>>> a[np.arange(a.shape[0]), amin_index]
array([-0.0069325 ,  0.04268358, -0.00128002, -0.01185333, -0.00389487])

这是基本索引(而不是高级索引)，因此返回的数组实际上是 a 的 View ，而不是内存中的新数组。