如何用(N-1)维数组高效访问N维数组？

使用 (N-1) 维数组访问 N 维数组

给定一个 N 维数组 a 和一个 (N- 1）维数组idx，一个常见的任务是访问idx中索引指定的元素。这对于执行查找最大值或检索特定值等操作非常有用。

使用高级索引的优雅解决方案

一个优雅的解决方案涉及使用 NumPy 的 ogrid 函数进行高级索引:

<code class="python">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]</code>

这将创建一个索引网格，并使用它来索引 a 和 b，从而生成包含相应值的数组。

函数的一般情况

对于适用于任何指定轴的更通用的解决方案，我们可以定义一个函数：

<code class="python">def argmax_to_max(arr, argmax, 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)]</code>

此函数接受一个数组、其沿指定轴的 argmax 以及轴本身。然后它构造一个网格并使用它来提取相应的元素。

使用自定义函数简化索引

为了进一步简化索引过程，我们可以创建一个辅助函数生成索引网格：

<code class="python">def all_idx(idx, axis):
    grid = np.ogrid[tuple(map(slice, idx.shape))]
    grid.insert(axis, idx)
    return tuple(grid)</code>

此函数返回一个索引元组，可直接用于索引输入数组：

<code class="python">axis = 0
a_max_values = a[all_idx(idx, axis=axis)]
b_max_values = b[all_idx(idx, axis=axis)]</code>

以上是如何用(N-1)维数组高效访问N维数组？的详细内容。更多信息请关注PHP中文网其他相关文章！