Question

我正在尝试为3D ndarray找到沿特定轴（0）的最大值的索引，然后使用这些索引切出这些值（以及来自第二个并行数组的相应值）。例如，

> a = np.random.randint(10, 100, 24).reshape(2, 3, 4)
> print(a)
array([[[94, 22, 96, 44],
        [11, 85, 39, 85],
        [58, 43, 48, 84]],

       [[84, 58, 51, 30],
        [74, 89, 90, 11],
        [90, 54, 94, 20]]])

现在，我对能够给出第0轴上最大值的指数感兴趣，即

> a[inds]
array([[94, 58, 96, 44],
       [74, 89, 90, 85,],
       [90, 54, 94, 84,]])

使用a.argmax()给出哪个第0轴索引是最大值，即

> a.argmax(axis=0)
array([[0, 1, 0, 0],
       [1, 1, 1, 0],
       [1, 1, 1, 0]])

但这不适用于切片...

Answer 1

在沿其余轴/ dims创建范围数组时使用advanced-indexing，如下所示 -

m,n = a.shape[1:]
Y,Z = np.ogrid[:m,:n]
Y_max_axis0 = a[inds,Y,Z]

示例运行 -

In [15]: a
Out[15]: 
array([[[94, 22, 96, 44],
        [11, 85, 39, 85],
        [58, 43, 48, 84]],

       [[84, 58, 51, 30],
        [74, 89, 90, 11],
        [90, 54, 94, 20]]])

In [16]: inds = a.argmax(axis=0)

In [17]: m,n = a.shape[1:]
    ...: Y,Z = np.ogrid[:m,:n]
    ...: Y_max_axis0 = a[inds,Y,Z]
    ...: 

In [18]: Y_max_axis0
Out[18]: 
array([[94, 58, 96, 44],
       [74, 89, 90, 85],
       [90, 54, 94, 84]])

创建这些范围数组然后编制索引的更明确的方法 -

In [19]: a[inds,np.arange(a.shape[1])[:,None], np.arange(a.shape[2])]
Out[19]: 
array([[94, 58, 96, 44],
       [74, 89, 90, 85],
       [90, 54, 94, 84]])

切片以在ndarray中沿轴提取最大值

1 个答案: