具有(n-1)d数组的索引n维数组 [英] Index n dimensional array with (n-1) d array

问题描述

如虚拟示例中那样，沿给定维访问具有(n-1)维数组的n维数组的最优雅方法是什么

What is the most elegant way to access an n dimensional array with an (n-1) dimensional array along a given dimension as in the dummy example

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指定的值?

How can I access now with idx a to get the maxima in a as if I had used a.max(axis=0)? or how to retrieve the values specified by idx in b?

我考虑过使用np.meshgrid，但是我认为这太过分了.请注意，尺寸axis可以是任何有用的轴(0,1,2)，并且事先未知.有没有一种优雅的方法可以做到这一点?

I thought about using np.meshgrid but I think it is an overkill. Note that the dimension axis can be any usefull axis (0,1,2) and is not known in advance. Is there an elegant way to do this?

推荐答案

使用

---

对于一般情况:

---

For the general case:

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)]

不幸的是，比这种自然的操作要尴尬得多.

Quite a bit more awkward than such a natural operation should be, unfortunately.

要用(n-1) dim数组索引n dim数组，我们可以简化一下，以便为我们提供所有轴的索引网格，就像这样-

For indexing a n dim array with a (n-1) dim array, we could simplify it a bit to give us the grid of indices for all axes, like so -

def all_idx(idx, axis):
    grid = np.ogrid[tuple(map(slice, idx.shape))]
    grid.insert(axis, idx)
    return tuple(grid)

因此，用它来索引输入数组-

Hence, use it to index into input arrays -

axis = 0
a_max_values = a[all_idx(idx, axis=axis)]
b_max_values = b[all_idx(idx, axis=axis)]

这篇关于具有(n-1)d数组的索引n维数组的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！