多维 np.argmax？

Question

我有一个形状为 (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 不支持多个轴作为参数。

无论如何，有没有解决方案可以得到这个结果？

Answer 1

看来您正在寻求为切片的每个扁平版本获取二维（行，列）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)