irascibility：argmax多轴无循环

irascibility：argmax多轴无循环

你可以做这样的事情-

# 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 Amaxpos_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 pre     ...: 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维数组

我们可以泛化来解决n-dim数组以获取

argmax

最后的

N

轴，并结合类似以下内容：

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

。