时间戳

首页 > 解决方案 > 如何获取 2-dim numpy 数组对的 topK 项？

pandas - 如何获取 2-dim numpy 数组对的 topK 项？

问题描述

我有 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 项，按分数第二行排序

我想我可以通过排序来实现：

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

但是 numpy 或 pandas 库有没有更优雅的方法？

标签： pandasnumpy

解决方案

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

推荐阅读