argmax usage

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))#输出每列最大元素所在下标
print(np.argmax(a, axis=1))#输出每行最大元素所在下标