详解numpy的argmax的具体使用

argmax函数基础概念

从最简单的例子出发

假设存在一个数组a = [3, 1, 2, 4, 6, 1]，我们需要计算该数组中最大数值的索引。对于初学者而言，这是一个基础问题。最直观的思路是先假设第0个元素为最大值，记录其索引，随后将其与后续元素逐一比较，若发现更大的数则更新索引。实现代码如下：

a = [3, 1, 2, 4, 6, 1]
maxindex = 0
i = 0
for tmp in a:
  if tmp > a[maxindex]:
    maxindex = i
  i += 1
print(maxindex)

虽然该问题看似简单，但它有助于我们理解argmax函数的基本逻辑。

解释

一维数组中的argmax应用

我们继续以一维数组为例，参考以下代码片段。

import numpy as np
a = np.array([3, 1, 2, 4, 6, 1])
print(np.argmax(a))

二维数组的axis参数解析

argmax函数的作用是返回数组中最大值的索引位置。该函数包含一个可选参数axis，默认值为0，用于指定沿哪一维寻找最大值。接下来我们观察二维数组的情况。

import numpy as np
a = np.array([[1, 5, 5, 2],
       [9, 6, 2, 8],
       [3, 7, 9, 1]])
print(np.argmax(a, axis=0))

为了便于描述，将上述二维数组记为a。调用np.argmax(a, axis=0)时，其含义是对每一列j（j=0,1,2,3），在a[0][j]、a[1][j]和a[2][j]中找出最大值对应的行索引。初始时，最大值索引假定为(0,0,0,0)。将第一行a[0][j]与第二行a[1][j]进行比较：由于9大于1、6大于5、8大于2，最大值索引更新为(1,1,0,1)。随后将更新后的结果与第三行a[2][j]比较：因为7大于6、9大于5，索引进一步更新为(1,2,2,1)。具体输出结果如下所示。

import numpy as np
a = np.array([[1, 5, 5, 2],
       [9, 6, 2, 8],
       [3, 7, 9, 1]])
print(np.argmax(a, axis=1))

调用np.argmax(a, axis=1)时，其含义是对每一行i（i=0,1,2），在a[i][0]、a[i][1]、a[i][2]和a[i][3]中找出最大值对应的列索引。初始时，假设a[i][0]即为最大值，对应索引为(0,0,0)。将a[i][0]与a[i][1]比较：由于5大于1、7大于3，索引更新为(1,0,1)。接着与a[i][2]比较：由于9大于7，索引更新为(1,0,2)。最后与a[i][3]比较，无需更新，最终结果为(1,0,2)。

import numpy as np
a = np.array([
       [
         [1, 5, 5, 2],
         [9, -6, 2, 8],
         [-3, 7, -9, 1]
       ],

       [
         [-1, 5, -5, 2],
         [9, 6, 2, 8],
         [3, 7, 9, 1]
       ]
      ])
print(np.argmax(a, axis=0))

三维数组的多维索引查找

接下来分析三维数组的情况。当axis=0时，np.argmax(a, axis=0)的含义是在a[0][j][k]与a[1][j][k]之间（j=0,1,2；k=0,1,2,3）寻找最大值索引。初始索引设为((0,0,0,0),(0,0,0,0),(0,0,0,0))。比较对应位置元素：6大于-6、3大于-3、9大于-9，因此更新这些位置的索引，将原索引更新为((0,0,0,0),(0,1,0,0),(1,0,1,0))。随后分析axis=1的情况。

import numpy as np
a = np.array([
       [
         [1, 5, 5, 2],
         [9, -6, 2, 8],
         [-3, 7, -9, 1]
       ],

       [
         [-1, 5, -5, 2],
         [9, 6, 2, 8],
         [3, 7, 9, 1]
       ]
      ])
print(np.argmax(a, axis=1))

当axis=1时，np.argmax(a, axis=1)的含义是在a[i][0][k]与a[i][1][k]之间（i=0,1；k=0,1,2,3）寻找最大值索引。初始索引设为((0,0,0,0),(0,0,0,0))。比较对应位置元素：9大于1、8大于2、9大于-1、6大于5、2大于-5、8大于2，因此更新这些位置的索引，得到((1,0,0,1),(1,1,1,1))。此时最大值对应的数组元素为((9,5,5,8),(9,6,2,8))。接着将当前最大值数组与a[i][2][k]比较：7大于5、7大于6、9大于2，更新相应位置索引，最终得到((1,2,0,1),(1,2,2,1))。axis=2的情况处理逻辑类似。

参考资料

numpy官方文档

以上就是本文的全部内容，希望对大家的学习有所帮助，也希望大家多多支持本站。