在虚拟示例中，沿给定维度访问具有 (n-1) 维数组的 n 维数组的最优雅方法是什么
a = np.random.random_sample((3,4,4))
b = np.random.random_sample((3,4,4))
idx = np.argmax(a, axis=0)
我现在如何使用idx a 访问以获得a 中的最大值，就像我使用a.max(axis=0) 一样？或者如何在b中检索idx指定的值？
我曾考虑过使用np.meshgrid，但我认为这太过分了。请注意，尺寸axis 可以是任何有用的轴 (0,1,2)，并且事先不知道。有没有优雅的方法来做到这一点？
【问题讨论】：
m,n = a.shape[1:]
I,J = np.ogrid[:m,:n]
a_max_values = a[idx, I, J]
b_max_values = b[idx, I, J]
对于一般情况：
def argmax_to_max(arr, argmax, axis):
"""argmax_to_max(arr, arr.argmax(axis), axis) == arr.max(axis)"""
new_shape = list(arr.shape)
del new_shape[axis]
grid = np.ogrid[tuple(map(slice, new_shape))]
grid.insert(axis, argmax)
return arr[tuple(grid)]
不幸的是，这比这种自然的操作要尴尬得多。
为了用(n-1) dim 数组索引n dim 数组，我们可以稍微简化一下，为我们提供所有轴的索引网格，就像这样 -
def all_idx(idx, axis):
grid = np.ogrid[tuple(map(slice, idx.shape))]
grid.insert(axis, idx)
return tuple(grid)
因此，使用它来索引输入数组 -
axis = 0
a_max_values = a[all_idx(idx, axis=axis)]
b_max_values = b[all_idx(idx, axis=axis)]
【讨论】：
all_idx 真的很优雅。喜欢你的帖子，@Divakar！
a 是一维数组，那么使用元组，我们将得到一个标量，它复制 .max() 的行为。但是如果没有元组，我们将得到一个包含一个元素的数组。所以，也许这就是我认为将其保留为元组的原因之一。
all_idx 很好。我没有意识到使用 argmax 输出的形状而不是原始数组的形状会简化事情。至于元组而不是列表，高级索引语义使用元组更清晰。由于在某些条件下将列表转换为元组的向后兼容性处理（未完全正确记录）位，该列表在这种情况下的行为相同。在 NumPy 索引中将列表视为元组以及将列表视为数组时可能会令人惊讶，因此我更喜欢显式创建元组。
idx 是标量的情况是我什至没有考虑过的，在这种情况下，不会触发向后兼容性处理，如果我们保留一个列表而不是一个元组，则会出现错误的结果。如此证明，元组的行为更加一致且更易于预测。
take_along_axis 已被添加以简化此操作。还有一个put