索引n维数组与(n-1)d数组

Question

使用沿给定维度的(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)？或者如何检索idxin中指定的值b？

我想过使用,np.meshgrid但我认为这是一种矫枉过正.请注意,尺寸axis可以是任何有用的轴(0,1,2),并且事先不知道.有一种优雅的方式来做到这一点？

Answer 1

利用advanced-indexing-

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 dim数组索引(n-1) 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)]