用 (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?

推荐答案

利用 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]

<小时>

对于一般情况:

---

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屋！