Как получить элементы topK для пар массивов 2-dim numpy?

Question

У меня есть 2 парных 2-dim numpy массива (скажем, метки и оценки)

labels = np.array([['a','b','c','d'],
                    ['a1','b1','c1','d1']])
scores = np.array([[0.1, 0.2, 0.3,0.4],
                  [1,2,3,4]])

Я хочу получить из них лучшие k элементов, отсортированные по второй строке

I думаю, я могу добиться этого, отсортировав:

[scores[i][1], scores[i][0], labels[i][0], labels[i][1]]

Но есть ли более элегантный способ с библиотекой numpy или pandas?

Answer 1

numpy .argsort должен это сделать

import numpy as np

labels = np.array([['a','b','c','d'],
                    ['a1','b1','c1','d1']])
scores = np.array([[0.1, 0.2, 0.3,0.4],
                  [1,2,3,4]])

k = 2 # number of "top items"
idx = np.argsort(scores[-1])[-k:] # get the indices of top values

topkScores = scores[:,idx].T # you can remove .T if you wish to get the score pairs in columns