如何对二维numpy矩阵元素从大到小排序并保留对应的ij索引

核心思路 #

将二维矩阵的元素与对应的原始行列索引绑定后，做全局降序排序即可，方案天然支持整数、浮点数类型矩阵，也可兼容对称矩阵的补全需求，完全适配10x10到40x40的矩阵尺寸。

完整实现代码 #

import numpy as np

# 输入示例矩阵
A = np.array([[1, 2, 3],
              [0, 4, 5],
              [0, 0, 6]])

# --- 可选：如果是仅存上三角的对称矩阵，先补全全矩阵，不需要可注释此行 ---
# A = A + A.T - np.diag(np.diag(A))

# 获取所有位置的行列索引
rows, cols = np.indices(A.shape)
# 组合为 (元素值, 行索引, 列索引) 的三元组
item_index_pairs = list(zip(A.flatten(), rows.flatten(), cols.flatten()))
# 按元素值降序全局排序
sorted_res = sorted(item_index_pairs, key=lambda x: x[0], reverse=True)

# 按要求格式输出
for val, i, j in sorted_res:
    print(f"{val}, ({i},{j})")

运行输出 #

6, (2,2)
5, (1,2)
4, (1,1)
3, (0,2)
2, (0,1)
1, (0,0)

关于np.argsort的问题说明 #

你之前得到逐行独立排序的结果，是因为np.argsort默认沿数组最后一维（行维度）单独排序，如果想用np.argsort实现全局排序，可以先把矩阵打平为一维数组再处理，纯NumPy的实现参考如下：

flatten_A = A.flatten()
# 得到全局降序排序的一维索引
sorted_flat_idx = np.argsort(-flatten_A)
# 一维索引转换为原始二维ij索引
sorted_i, sorted_j = np.unravel_index(sorted_flat_idx, A.shape)
# 组合结果
sorted_res = list(zip(flatten_A[sorted_flat_idx], sorted_i, sorted_j))

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