python(np.argpartition())输出 前n个最大/小值所对应的索引

利用python numpy标准函数np.argpartition()，输出 前n个最大/小值所对应的索引（输出的索引是无序的）

需要注意：

1、输出索引所对应的数值是乱序的。这种方法复杂度要低

import numpy as np


def topk_partition(matrix, k, top_max=True):
    '''
    输出matrix矩阵的最大/小前k个值的index，无序
    所对应的top k的值：matrix[index[i][0], index[i][1]]
                        matrix[index[i]]
    :param matrix: 二维或一维
    :param k:
    :param top_max: True:最大值的前k个值，无序；False:最小值的前k个值，无序；
    :return:
    '''
    flag_2d = False
    if len(matrix.shape) == 2:
        flag_2d = True
        matrix1 = matrix.reshape((matrix.shape[0] * matrix.shape[1]))
    else:
        matrix1 = matrix.copy()
    if top_max: # matrix前k个最大值 所对应的index
        index = np.argpartition(matrix1, -k)[-k:len(matrix1)]
    else:   # matrix前k个最小值 所对应的index
        index = np.argpartition(matrix1, k-1)[0:k]

    # 解析index
    if flag_2d:
        index_2d = []
        for i in index:
            row = i // matrix.shape[1]
            col = i % matrix.shape[1]
            index_2d.append([row,col])
        index = index_2d

    return index

if __name__ == '__main__':
    # dists= np.array([ 6,0,2,1,5,1.2,4,22])
    dists = np.random.randint(0, 10, size=(3, 4))
    print(dists)
    index = topk_partition(dists, 4, top_max=False)
    print(index)
    # print(dists[index])