如何快速获取二维数组每行的N个最大值对应索引？

当处理K×M（M>10000）的大规模相似度矩阵时，原方案用np.argsort做全排序效率极低，以下是几种实用的优化方法：

1. 用np.argpartition替代全排序 #

np.argsort会对每行所有元素排序，时间复杂度为O(KM log M)；而np.argpartition只需找到前N大元素的位置，时间复杂度为O(KM)，之后仅对这N个元素局部排序即可：

import numpy as np

def get_args_of_best_score_fast(score_matrix, N):
    # 获取每行前N大元素的索引（未排序）
    partitioned_indices = np.argpartition(score_matrix, -N, axis=1)[:, -N:]
    # 对每行的N个元素按相似度从高到低排序
    for i in range(score_matrix.shape[0]):
        row_scores = score_matrix[i, partitioned_indices[i]]
        # 按负分排序等价于按原分数降序
        sorted_sub_idx = np.argsort(-row_scores)
        partitioned_indices[i] = partitioned_indices[i][sorted_sub_idx]
    return partitioned_indices

2. 直接用向量检索库（如Faiss）跳过相似度矩阵计算 #

如果相似度矩阵是由向量内积/余弦相似度推导而来，无需生成完整的K×M矩阵，直接用Faiss做Top-N检索，底层为C++优化实现，效率远超纯Python方案：

import faiss

def faiss_top_n_search(query_vectors, db_vectors, N, metric=faiss.METRIC_INNER_PRODUCT):
    # 根据度量方式构建索引，L2距离对应欧氏距离，IP对应内积（余弦相似度可归一化后用IP）
    if metric == faiss.METRIC_L2:
        index = faiss.IndexFlatL2(db_vectors.shape[1])
    else:
        index = faiss.IndexFlatIP(db_vectors.shape[1])
    index.add(db_vectors)
    # 检索每个查询向量的Top-N匹配结果
    _, top_n_indices = index.search(query_vectors, N)
    return top_n_indices

3. 用Numba JIT加速循环 #

对纯numpy操作的循环部分做JIT编译，将Python代码转为机器码执行，大幅提升循环效率：

from numba import jit
import numpy as np

@jit(nopython=True)
def get_args_of_best_score_numba(score_matrix, N):
    K, M = score_matrix.shape
    result = np.zeros((K, N), dtype=np.int64)
    for i in range(K):
        # 对当前行降序排序后取前N个索引
        sorted_idx = np.argsort(-score_matrix[i])[:N]
        result[i] = sorted_idx
    return result

4. 分块处理优化内存与缓存命中率 #

将大矩阵拆分为小块处理，降低内存占用的同时，利用CPU缓存提升计算速度：

import numpy as np

def get_args_of_best_score_chunked(score_matrix, N, chunk_size=1000):
    K, M = score_matrix.shape
    result = np.zeros((K, N), dtype=np.int64)
    # 按块处理每行数据
    for start in range(0, K, chunk_size):
        end = min(start + chunk_size, K)
        chunk = score_matrix[start:end]
        # 对块内数据做partition操作
        partitioned = np.argpartition(chunk, -N, axis=1)[:, -N:]
        # 局部排序块内的Top-N结果
        for i in range(end - start):
            row_scores = chunk[i, partitioned[i]]
            sorted_sub_idx = np.argsort(-row_scores)
            partitioned[i] = partitioned[i][sorted_sub_idx]
        result[start:end] = partitioned
    return result

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