numpy多维数组取最大值的索引及前topk最大值的索引

主要介绍numpy的 argmax，argsort函数，以及heapq库

# 返回最大值的索引
a = np.array([[6.8284258e-09, 5.7629350e-09, 7.4942248e-09,9.5846470e-11,
        5.4229814e-14, 8.6114553e-15],
       [1.1742729e-09, 9.2132124e-10, 1.2312825e-09,2.3951413e-04,
        5.2452291e-09, 7.0311121e-09]])
d=[c.argmax() for c in a]
print(d)
>>>[2, 3]


# argsort
# np.argsort(a, axis=-1, kind='quicksort', order=None)
Examples
--------
One dimensional array:

>>> x = np.array([3, 1, 2])
>>> np.argsort(x)
array([1, 2, 0])

Two-dimensional array:

>>> x = np.array([[0, 3], [2, 2]])
>>> x
array([[0, 3],
       [2, 2]])

>>> np.argsort(x, axis=0)  # sorts along first axis (down)
array([[0, 1],
       [1, 0]])

>>> np.argsort(x, axis=1)  # sorts along last axis (across)
array([[0, 1],
       [0, 1]])

Indices of the sorted elements of a N-dimensional arr