我有一個(gè)數(shù)組:a = np.array([[22,11,44,33,66], [22,11,2,1,66], [1,11,44,22,4], [22,11,88,99,66]])作為輸出,我想要一個(gè)包含最大 3 個(gè)值的索引的數(shù)組作為二維數(shù)組。例如對(duì)于上面的數(shù)組輸出將是:array([[4,2,3], [4,0,1], [2,3,1], [3,2,4]])
1 回答

慕森卡
TA貢獻(xiàn)1806條經(jīng)驗(yàn) 獲得超8個(gè)贊
要獲取k數(shù)組的頂部元素,請(qǐng)對(duì)其進(jìn)行分區(qū)。由于分區(qū)通常為您提供k最低元素,請(qǐng)使用反向索引:
k = 3
top = np.argpartition(a, -k, axis=1)[:, -k:]
如果您需要按降序?qū)λ饕M(jìn)行排序,請(qǐng)np.argsort與結(jié)果一起使用:
rows = np.arange(a.shape[0])[:, None]
s = np.argsort(a[rows, top], axis=1)[:, ::-1]
top = top[rows, s]
rows當(dāng)您使用top和進(jìn)行花哨的索引時(shí),有必要確保正確選擇所有索引s。每行的索引必須顛倒 ( [:, ::-1]) 以獲得升序。
添加回答
舉報(bào)
0/150
提交
取消