Python numpy.argmax函数方法的使用

numpy.argmax

numpy.argmax(a, axis=None, out=None)     [source]

沿着一个轴返回最大值的索引。

Notes

如果最大值多次出现，则返回与第一次出现对应的索引。

例子

1）二维数组的最大值索引

import numpy as np

# 创建一个二维数组
a = np.arange(6).reshape(2, 3) + 10
print("Array a:")
print(a)

# 获取整个数组中最大值的索引
max_index = np.argmax(a)
print(f"max_index: {max_index}")

# 沿着第0轴（列）计算最大值的索引
max_index_axis0 = np.argmax(a, axis=0)
print(f"max_index_axis0: {max_index_axis0}")

# 沿着第1轴（行）计算最大值的索引
max_index_axis1 = np.argmax(a, axis=1)
print(f"max_index_axis1: {max_index_axis1}")
import numpy as np

# 创建一个二维数组
a = np.arange(6).reshape(2, 3) + 10
print("Array a:")
print(a)

# 获取整个数组中最大值的索引
max_index = np.argmax(a)
print(f"max_index: {max_index}")

# 沿着第0轴（列）计算最大值的索引
max_index_axis0 = np.argmax(a, axis=0)
print(f"max_index_axis0 : {max_index_axis0}")

# 沿着第1轴（行）计算最大值的索引
max_index_axis1 = np.argmax(a, axis=1)
print(f"max_index_axis1: {max_index_axis1}")

2）获取N维数组最大值的多维索引

import numpy as np

# 创建一个二维数组
a = np.arange(6).reshape(2, 3) + 10
print("Array a:")
print(a)

# 使用 np.unravel_index 获取多维数组中最大值的索引
ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
print(f"ind: {ind}")

# 获取最大值
print(f"Maximum value: {a[ind]}")

3）使用 np.argmax() 获取最大值索引并使用 np.take_along_axis() 获取最大值

import numpy as np

# 创建一个二维数组
x = np.array([[4, 2, 3], [1, 0, 3]])

# 获取最大值所在索引的数组
index_array = np.argmax(x, axis=-1)
print(f"index_array: {index_array}")

# 使用 np.take_along_axis 获取最大值
max_values = np.take_along_axis(x,
np.expand_dims(index_array, axis=-1),
axis=-1)
print(f"max_values:\n{max_values}")

# 移除多余的维度，得到最终的最大值
max_values_squeezed = max_values.squeeze(axis=-1)
print(f"max_values_squeezed: {max_values_squeezed}")