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

【问题标题】：Slicing to extract maximum values along an axis, in an ndarray切片以在 ndarray 中沿轴提取最大值
【发布时间】：2017-06-07 01:56:27
【问题描述】：

我正在尝试查找 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]]])

现在，我对在第零轴上给出最大值的索引感兴趣，即

> 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]])

但这不适用于切片...

【问题讨论】：

标签： arrays numpy math slice

【解决方案1】：

使用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]])

【讨论】：

非常感谢！我需要对我的高级索引进行（很多）复习......但是有一个简单的原因为什么不需要将第三轴的索引（np.arange(a.shape[2])）广播到一个额外的维度（即np.arange(a.shape[2])[:, None] )?
@DilithiumMatrix 不，我们不需要。这就是np.arange(a.shape[2]) 中显示的explicit-way 也可以工作的原因。但是使用由 np.ogrid（如前所示）给出的附加暗淡，我们有 np.arange(a.shape[2])[None] 也不会受到伤害。
@DilithiumMatrix inds 将具有 (m,n) 的暗淡。我们需要扩展np.arange(m)，因为它沿着inds 的第一轴，而np.arange(n) 沿着inds 的第二轴。所以，我们需要扩展第一个范围数组而不是第二个，因为第二个在索引时会自动沿着inds 的第一个轴扩展。