Numpy:对于一维中的每个元素,找到子数组最大值的坐标

Numpy: for each element in one dimension, find coordinates of maximum of sub-array

我已经看到这个问题的变体被问过几次,但到目前为止还没有看到任何能触及这个一般案例核心的答案。我有一个形状为 [a, b, c, ...] 的 n 维数组。对于某些维度x,我想查看每个子数组并找到最大值的坐标。

例如,说b = 2,这就是我感兴趣的维度。我想要[:, 0, :, ...][:, 1, :, ...]中最大值的坐标,形式为[=18] =]等

我尝试通过将输入矩阵重塑为二维数组 [b, a*c*d*...]、沿 轴 0 使用 argmax 并解开索引,但输出坐标不会在我的数据集中给出最大值。在这种情况下,n = 3 并且我对 轴 1 感兴趣。

shape = gains_3d.shape
idx = gains_3d.reshape(shape[1], -1) 
idx = idx.argmax(axis = 1)
a1, a2 = np.unravel_index(idx, [shape[0], shape[2]])

显然我可以使用循环,但这不是很 pythonic。

举个具体的例子,我随机生成了一个4x2x3的数组。我对轴 1 感兴趣,所以输出应该是两个长度为 2 的数组。

testarray = np.array([[[0.17028444, 0.38504759, 0.64852725],
        [0.8344524 , 0.54964746, 0.86628204]],

       [[0.77089997, 0.25876277, 0.45092835],
        [0.6119848 , 0.10096425, 0.627054  ]],

       [[0.8466859 , 0.82011746, 0.51123959],
        [0.26681694, 0.12952723, 0.94956865]],

       [[0.28123628, 0.30465068, 0.29498136],
        [0.6624998 , 0.42748154, 0.83362323]]])

testarray[:,0,:]

array([[0.17028444, 0.38504759, 0.64852725],
       [0.77089997, 0.25876277, 0.45092835],
       [0.8466859 , 0.82011746, 0.51123959],
       [0.28123628, 0.30465068, 0.29498136]])

,所以第一个输出数组的第一个元素会是2,另一个的第一个元素会是0,指向0.8466859。两个矩阵的第二个元素将是2和2,指向testarray[:,1,:]

的0.94956865

让我们首先尝试清楚地了解您要做什么:

示例 3d 数组:

In [136]: arr = np.random.randint(0,10,(2,3,4))
In [137]: arr
Out[137]: 
array([[[1, 7, 6, 2],
        [1, 5, 7, 1],
        [2, 2, 5, *6*]],

       [[*9*, 1, 2, 9],
        [2, *9*, 3, 9],
        [0, 2, 0, 6]]])

经过一番摆弄后,我想出了这个迭代,显示每个中间维度的坐标和最大值

In [151]: [(i,np.unravel_index(np.argmax(arr[:,i,:]),(2,4)),np.max(arr[:,i,:])) for i in range
     ...: (3)]
Out[151]: [(0, (1, 0), 9), (1, (1, 1), 9), (2, (0, 3), 6)]

我可以将 unravel 移到迭代之外:

In [153]: np.unravel_index([np.argmax(arr[:,i,:]) for i in range(3)],(2,4))
Out[153]: (array([1, 1, 0]), array([0, 1, 3]))

您的重塑方法确实避免了这个循环:

In [154]: arr1 = arr.transpose(1,0,2)  # move our axis first
In [155]: arr1 = arr1.reshape(3,-1)    
In [156]: arr1
Out[156]: 
array([[1, 7, 6, 2, 9, 1, 2, 9],
       [1, 5, 7, 1, 2, 9, 3, 9],
       [2, 2, 5, 6, 0, 2, 0, 6]])

In [158]: np.argmax(arr1,axis=1)
Out[158]: array([4, 5, 3])
In [159]: np.unravel_index(_,(2,4))
Out[159]: (array([1, 1, 0]), array([0, 1, 3]))

maxargmax 仅采用一个轴值,而您想要的相当于沿除一个轴以外的所有轴取最大值。一些 ufunc 采用轴元组,但这些不采用。转置和重塑可能是唯一的方法。

In [163]: np.max(arr1,axis=1)
Out[163]: array([9, 9, 6])