Numpy:多个轴上的 argmax 没有循环 [英] Numpy: argmax over multiple axes without loop
问题描述
我有一个 N 维数组(命名为 A).对于 A 的第一轴的每一行,我想获取沿 A 其他轴的最大值的坐标.然后我将返回一个二维数组,其中包含第一轴每一行的最大值的坐标A.
I have a N-dimensional array (Named A). For each row of the first axis of A, I want to obtain the coordinates of the maximum value along the other axes of A. Then I would return a 2-dimensional array with the coordinates of the maximum value for each row of the first axis of A.
我已经使用循环解决了我的问题,但我想知道是否有更有效的方法来做到这一点.我目前的解决方案(例如数组 A)如下:
I already solved my problem using a loop, but I was wondering whether there is a more efficient way of doing this. My current solution (for an example array A) is as follows:
import numpy as np
A=np.reshape(np.concatenate((np.arange(0,12),np.arange(0,-4,-1))),(4,2,2))
maxpos=np.empty(shape=(4,2))
for n in range(0, 4):
maxpos[n,:]=np.unravel_index(np.argmax(A[n,:,:]), A[n,:,:].shape)
在这里,我们会:
A:
[[[ 0 1]
[ 2 3]]
[[ 4 5]
[ 6 7]]
[[ 8 9]
[10 11]]
[[ 0 -1]
[-2 -3]]]
maxpos:
[[ 1. 1.]
[ 1. 1.]
[ 1. 1.]
[ 0. 0.]]
如果有多个最大化器,我不介意选择哪个.
If there are multiple maximizers, I don't mind which is chosen.
我尝试使用 np.apply_over_axes
,但我没有设法让它返回我想要的结果.
I have tried to use np.apply_over_axes
, but I haven't managed to make it return the outcome I want.
推荐答案
你可以这样做 -
# Reshape input array to a 2D array with rows being kept as with original array.
# Then, get idnices of max values along the columns.
max_idx = A.reshape(A.shape[0],-1).argmax(1)
# Get unravel indices corresponding to original shape of A
maxpos_vect = np.column_stack(np.unravel_index(max_idx, A[0,:,:].shape))
样品运行 -
In [214]: # Input array
...: A = np.random.rand(5,4,3,7,8)
In [215]: # Setup output array and use original loopy code
...: maxpos=np.empty(shape=(5,4)) # 4 because ndims in A is 5
...: for n in range(0, 5):
...: maxpos[n,:]=np.unravel_index(np.argmax(A[n,:,:,:,:]), A[n,:,:,:,:].shape)
...:
In [216]: # Proposed approach
...: max_idx = A.reshape(A.shape[0],-1).argmax(1)
...: maxpos_vect = np.column_stack(np.unravel_index(max_idx, A[0,:,:].shape))
...:
In [219]: # Verify results
...: np.array_equal(maxpos.astype(int),maxpos_vect)
Out[219]: True
推广到 n-dim 数组
我们可以概括来解决 n-dim 数组以获得最后 N
个轴的 argmax
结合类似这样的东西 -
Generalize to n-dim array
We could generalize to solve for n-dim array to get argmax
for last N
axes combined with something like this -
def argmax_lastNaxes(A, N):
s = A.shape
new_shp = s[:-N] + (np.prod(s[-N:]),)
max_idx = A.reshape(new_shp).argmax(-1)
return np.unravel_index(max_idx, s[-N:])
结果将是一个索引数组的元组.如果您需要将最终输出作为数组,我们可以使用 np.stack
或 np.concatenate
.
The result would a tuple of arrays of indices. If you need the final output as an array, we can use np.stack
or np.concatenate
.
这篇关于Numpy:多个轴上的 argmax 没有循环的文章就介绍到这了,希望我们推荐的答案对大家有所帮助,也希望大家多多支持IT屋!