如何使用 argmin 的输出作为 Numpy 的索引

Question

我想在 rank-3 numpy 数组中找到沿给定轴的最小值的位置。我已经通过 np.argmin 获得了这些位置，但是我不确定如何将其“应用”到原始矩阵以获得实际的最小值。

例如：

import numpy as np

a = np.random.randn(10, 5, 2)
min_loc = a.argmin(axis = 0)   # this gives an array of shape (5, 2)

现在，问题是如何使用 min_loc 获得实际的最小值？我试过 a[min_loc]，它给了我一个形状 (5, 2, 5, 2)。这个形状的逻辑是什么？我如何使用这个辅助矩阵来获得形状 (5, 2)

的合理解

请注意，a.min(axis = 0) 不是我正在寻找的解决方案。我需要通过 argmin.

的解决方案

Answer 1

a[min_loc] 在第一个维度上执行 integer array indexing，即它将为 min_loc 中的每个索引选择 (5, 2) 形状的数组。由于 min_loc 本身是 (5, 2) 形状的，并且对于 min_loc 中的每个整数，它会选择另一个 (5, 2) 形状的数组。你最终得到一个 (5, 2, 5, 2) 数组。同样的原因 a[np.array([0, 3])] 具有 (2, 5, 2) 的形状，a[np.array([[0], [3]])] 具有 (2, 1, 5, 2) 的形状，因为您只提供第一维的索引。

对于您的用例，您不想为 min_loc 中的每个索引选择一个子数组，而是需要一个元素。例如，如果您有 min_loc = [[5, ...], ...]，第一个元素的完整索引应该是 5, 0, 0 而不是 5, :, :。这正是 advanced indexing 所做的。基本上通过提供一个整数数组作为每个维度的索引，您可以提取与特定位置对应的元素。您可以使用 np.indices:

从 (5, 2) 形状构建第 2 和第 3 维的索引

j, k = np.indices(min_loc.shape)
a[min_loc, j, k]

# [[-1.82762089 -0.80927253]
#  [-1.06147046 -1.70961507]
#  [-0.59913623 -1.10963768]
#  [-2.57382762 -0.77081778]
#  [-1.6918745  -1.99800825]]

其中 j, k 是第二维和第三维的坐标：

j
#[[0 0]
# [1 1]
# [2 2]
# [3 3]
# [4 4]]

k  
#[[0 1]
# [0 1]
# [0 1]
# [0 1]
# [0 1]]

或如@hpaulj 所评论，使用np.take_along_axis 方法：

np.take_along_axis(a, min_loc[None], axis=0)

# [[[-0.93515242 -2.29665325]
#   [-1.30864779 -1.483428  ]
#   [-1.24262879 -0.71030707]
#   [-1.40322789 -1.35580273]
#   [-2.10997209 -2.81922197]]]

如何使用 argmin 的输出作为 Numpy 的索引

How to use the output of argmin as index with Numpy

python

numpy

array-broadcasting