Python：从 3D 数组中的 argmin 广播二维索引数组

Question

我有一个具有不同值的 3D 网格，并将 argmin 沿 axis=2 应用于它，以获得第三维的最小值。 我现在如何提取实际值而不仅仅是最小值的索引？

import numpy as np
input = np.random.normal(size=(30,40,10))
minvals = np.argmin(input,axis=2)

foo  = input[minvals]

Minvals 为我提供沿轴 2 的索引，以获得预期的最小值。 我预计， foo 给了我一个具有实际最小值的二维数组，但 foo 现在是一个 4D 数组......

Answer 1

尝试take_along_axis与expand_dims混合。 有关详细信息，请参阅take_along_axis 教程：

import numpy as np

input = np.random.normal(size=(30,40,10))
minvals = np.argmin(input,axis=2)
foo = np.take_along_axis(input, np.expand_dims(minvals, axis=2), axis=2).reshape(30,40)

# Verification
foo_2 = np.min(input,axis=2)
print('Good result!' if np.allclose(foo,foo_2) else 'Bad results')

输出：

% python3 script.py
Good result!