基于 numpy.argmin 结果的切片

Question

让我们有一个形状等于 (36, 2, 400, 400) 的 numpy 数组 (float)。让我们说 400 x 400 代表一个图像。然后对于每个像素，我想找到两个值（第二维），它们是在第二维上取范数时，相对于第一维来说是最低的。所以我最终得到了一个形状数组 (2, 400, 400)。

使用np.argmin(np.linalg.norm(array, axis=1), axis=0)，我可以得到每个 2 x 400 x 400 像素的索引，这几乎是我想要的。但是现在我想使用这个数字在第一维中对原始数组进行切片，所以我留下了一个形状为 (2, 400, 400) 的数组。

我可以做的是遍历所有索引并逐像素构造结果，但我相信有一种更智能的方法。谁能用更聪明的方法帮助我？

根据要求提供的最小可重复示例，其中距离是数组：

shape = (400, 400)
centers = np.random.randint(400, size=(36, 2))

distances = np.array([np.indices(shape) - np.array(center)[:, None, None] for center in centers])
nearest_center_index = np.argmin(np.linalg.norm(distances, axis=1), axis=0)

print(distances.shape)
print(nearest_center_index.shape)
plt.imshow(nearest_center_index)

出来：

(36, 2, 400, 400)
(400, 400)

在 cmets 的帮助下，我能够得出一个有点难看的答案，这有助于我进一步理解问题。让我详细说明。可以做的是展平图像和 argmin 结果，然后在图像上使用带有 argmin 和索引的高级索引来生成结果。

flatten_indices = nearest_center_index.reshape(400**2)
image_indices = range(400**2)
results = distances.reshape(36, 2, 400**2)[flatten_indices, :, image_indices].reshape(400, 400, 2).swapaxes(0, 2)

但是，我认为很多情况下，您的索引都被塑造为维度的子集，并且具有包含另一个维度索引的值。我希望有一个通用的方法来分割它。

因此，让我们有一个 n 维数组，形状 = (x1, x2, ..., xn)，假设我们有一个数组表示一个维度的索引，例如 xi，它的形状是原始数组形状的子集，不包含 xi。然后我希望有一种方法可以对这个数组进行切片。

Answer 1

我要找的函数是numpy.take_along_axis()。

对于特定示例，唯一需要做的是确保nearest_center_index（argmin 的输出）与要切片的数组具有相同的维度。在示例中，这可以通过将keepdims=True 传递给norm 和argmin 来实现，然后可以直接将其用作numpy 函数的第二个参数。第三个参数应该是 xi 轴（在示例轴 0 中）。

不通过keepdims=True，按照确切的示例，可以通过以下方式实现所述目标：

result = np.take_along_axis(distances, nearest_center_index[None,None,:,:], 0)[0]