如何用NumPy按score矩阵对inference矩阵进行整体排序？

我有一个名为inference的矩阵，由存储itemId的数组组成，每个数组对应一位用户的推荐结果，示例如下：

inference[0] = array([ 1, 17,  0, 29, 33, 10, 23, 18,  4, 25, 37, 41, 19,  7, 45, 44, 28,
        5, 21, 30, 27,  6, 16, 32,  3, 46, 47, 11, 24, 35, 39, 15, 22, 31,
       43], dtype=int32)

还有一个名为score的矩阵，由对应索引itemId的预测分数数组组成，每个数组对应一位用户的itemId分数集合，示例如下：

score[0] = array([ 4.66423448e-01,  3.04435879e-01,  1.20756114e-01,  7.42338740e-03,
        1.00917931e-02,  3.40771784e-02,  2.95762312e-02,  4.64895252e-03,
       -4.86475747e-02, -5.37142403e-03, -2.96056704e-04, -3.23560827e-05,
       -2.89172482e-02, -3.72408911e-02, -6.24527574e-01, -1.06988378e-04,
       -1.80022987e-03, -3.40648238e-02, -2.07088395e-02, -2.53725616e-03,
       -2.20156523e-02, -3.26039633e-02, -5.12802875e-02, -1.61312032e-03,
       -1.99290374e-01, -1.46841628e-04, -8.44907165e-01, -1.73397407e-01,
       -3.57963537e-02, -1.43663881e-03, -1.67909664e-03, -5.75751424e-03,
       -2.39864983e-02, -3.77825587e-03, -9.72822814e-04])

即对于用户0，itemId#1的预测分数为4.66423448e-01，itemId#17的分数为3.04435879e-01，依此类推。

我希望按score矩阵对inference矩阵进行排序，示例排序后的数组如下：

sorted_matrix[0] = array([ 1, 17,  0, 10, 23, 33, 29, 18, 41, 44, 46, 37, 43, 35, 32, 39, 28,
       30, 31, 25, 15, 21, 27, 22, 19,  6,  5, 24,  7,  4, 16, 11,  3, 45,
       47], dtype=int32)

针对单个数组，我使用inference[0][np.argsort(-1 * (score))[0]]可以实现排序，但尝试对整个矩阵排序时，使用new_inference = inference[np.argsort(-score)]得到了嵌套矩阵，new_inference[0]变成了35x35的矩阵而非数组。请问我在使用np.argsort()时哪里出错了？

错误原因 #

当对二维score矩阵直接调用np.argsort(-score)时，argsort默认按**最后一维（列）**排序，返回一个与score形状相同的二维索引矩阵（比如你的数据是N×35，返回的索引矩阵也是N×35）。

此时用inference[np.argsort(-score)]索引会触发numpy的高级索引广播机制：索引矩阵的每一行会被用来取inference的所有行，最终生成一个N×N×35的三维数组，这就是new_inference[0]变成35×35矩阵的原因。

正确实现 #

要实现每个用户的推荐结果按对应分数降序排序，需要让索引与inference的行一一对应，具体步骤如下：

1. 对score按行排序（指定axis=1），得到每行的降序索引；
2. 生成匹配行维度的索引，确保每个用户的索引只取自己对应的推荐结果行。

代码示例：

import numpy as np

# 假设inference和score都是形状为(N, 35)的二维数组
sorted_indices = np.argsort(-score, axis=1)
# 生成行索引并扩展维度，与sorted_indices广播匹配
row_indices = np.arange(inference.shape[0])[:, None]
new_inference = inference[row_indices, sorted_indices]

代码解释 #

• np.argsort(-score, axis=1)：指定按每行（axis=1）进行降序排序，返回的sorted_indices是N×35的二维数组，每行对应一个用户的itemId排序索引。
• np.arange(inference.shape[0])[:, None]：生成形状为(N,1)的行索引数组，与sorted_indices的N×35形状广播后，形成N×35的索引矩阵，确保每个用户的索引仅作用于自己对应的inference行，最终得到与原矩阵形状一致的N×35排序结果。

内容的提问来源于stack exchange，提问作者DS Park