python - 在另一个数组中查找与沿轴的最小值对应的数组

Question

我有两个三维数组 a 和 b，我想找到 b 的二维子数组，其中 a 沿第三轴具有最小值的元素，即

a=n.random.rand(20).reshape((5,2,2))
b=n.arange(20).reshape((5,2,2))
c=n.argmin(a,2) #indices with minimal value of a
d=n.zeros_like(c) #the array I want
for i in range(5):
  for j in range(2):
     d[i,j] = b[i,j,c[i,j]] 

有没有办法不用双循环就可以得到这些值？

我知道这个答案: replace min value to another in numpy array但如果我想让它适用于我的 3D 阵列，我必须做很多 reshape 操作 - 我想知道是否有更简单的方法。

Answer 1

您可以使用 np.ogrid为其他维度创建网格:

x, y, z = arr.shape  # assuming your array is named "arr"
xyz = np.ogrid[0:x, 0:y] + [c]  # c is your second axis index (the argmin)
arr[xyz]

如果它不是最后一个轴，那么您可以简单地使用 insert 因为 ogrid 返回一个包含索引的普通 python 列表。