如何让Numpy argsort区分0值并输出对应-1标识？

全流程不做Python层逐元素循环，纯靠numpy向量化操作完成，完全匹配你需要的输出格式，适配大规模数组的运算性能要求。

1. 给0值设置专属排序占位 #

要保证非0值的原有排序顺序不受干扰，同时让原数组的0值全部沉到每行排序结果的末尾，需要把所有0值替换成一个远大于数组内所有非0有效值的数值。
针对你给的示例数组：

import numpy as np
a = np.array([[3, 0, 5, 0], 
              [8, 7, 6, 10], 
              [5, 4, 0, 10]])
# 把0值替换为比所有有效值都大的占位值，余弦相似度最大值为1，用np.inf不会和真实值冲突
a_masked = np.where(a == 0, np.inf, a)

2. 执行常规行级argsort #

对加了占位的数组做行排序，拿到排序后的索引：

sorted_idx = a_masked.argsort(axis=1)

这一步完成后，所有非0值的索引已经按数值从小到大排到每行的最前面，原数组里0值对应的索引，会因为占位值的高优先级自动落到每行的末尾。

3. 替换0值对应的索引为自定义标识 #

通过take_along_axis定位到排序结果中、对应原数组是0值的位置，把这些位置的索引替换为-1：

result = np.where(np.take_along_axis(a_masked, sorted_idx, axis=1) == np.inf, -1, sorted_idx)

最终输出完全符合预期效果：

array([[ 0,  2, -1, -1],
       [ 2,  1,  0,  3],
       [ 1,  0,  3, -1]])

实操注意事项 #

• 这套流程全是numpy底层的向量化运算，哪怕是十万行以上的余弦相似度矩阵，也能保持毫秒级的运算速度，不会出现大规模数组下的性能损耗
• 全0无效行的过滤逻辑可以直接复用结果：只要判断当前行是否全为-1，就能直接筛掉没有有效相似度的行，不需要额外计算掩码
• 之前尝试替换NaN、inf达不到预期的核心原因：替换值后没有对排序结果做二次标识替换，直接沿用了排序输出的原始索引，没有把0值对应的位置改成要求的特殊标识
• 版本适配优化：如果担心不同numpy版本对NaN、inf的排序逻辑有差异，可以提前计算数组的全局最大值，用全局最大值 + 1作为0值的占位值，保证占位值永远大于所有非0有效值，排序结果的一致性更强

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