时间戳

首页 > 解决方案 > 多维 np.argmax？

arrays - 多维 np.argmax？

问题描述

我有一个形状为 (n, n, g) 的 3D 数组，我需要每个 (n, n) argmax，即结果应该是两个长度为 g 的索引向量 (x, y)。

直观的解决方案是：

array = np.random.uniform(size=[5, 5, 1000])
np.argmax(array, axis=[0, 1])

但是，numpy 不支持多个轴作为参数。

有没有办法得到这个结果？

标签： arrayspython-3.xnumpy

解决方案

看来您正在为每个展平版本的切片获取二维（行，列）argmax 索引，2D其中前两个轴合并/组合。因此，第一个块将是array[:,:,0]等等，我们需要找到 argmax，该切片被展平并回溯到原始 2D 形状。因此，要解决它，我们可以简单地重塑以合并前两个轴，沿第一个轴获取 argmax，这是在重塑后合并的一个轴，并使用 回溯原始索引np.unravel_index，如下所示 -

array2D = array.reshape(-1,array.shape[-1])
r,c = np.unravel_index(array2D.argmax(0),array.shape[:2])

样品运行 -

In [29]: array = np.random.uniform(size=[5, 5, 1000])

In [30]: array2D = array.reshape(-1,array.shape[-1])

In [31]: r,c = np.unravel_index(array2D.argmax(0),array.shape[:2])

In [32]: len(r), len(c)
Out[32]: (1000, 1000)

让我们验证第一个2D切片的结果 -

In [33]: array[:,:,0]
Out[33]: 
array([[0.81590174, 0.17919069, 0.22717883, 0.67863625, 0.97390595],
       [0.82096447, 0.05894774, 0.86379174, 0.13494354, 0.10003756],
       [0.37243189, 0.33714008, 0.21165031, 0.35910642, 0.15163255],
       [0.1376776 , 0.86866599, 0.43602004, 0.85421372, 0.77805012],
       [0.10519547, 0.7422571 , 0.35632275, 0.24168307, 0.76882613]])

In [34]: array[:,:,0].argmax() 
Out[34]: 4  # flattened index for 0.97390595 at (0,4) in the 2D slice

In [36]: r[0],c[0]
Out[36]: (0, 4)

推荐阅读