如何高效从numpy矩阵中获取行最大值列且每列仅被选中一次

需求对齐 #

按行号从小到大依次处理，每行从未被占用的列中选择当前行值最大的列，每列仅能被选中一次，最终返回行-列匹配的元组列表。

实现方案 #

方案1：按行贪心实现（完全符合给出的规则） #

核心思路：先通过numpy矢量化操作预计算每行的列降序索引，再遍历行快速找到第一个可用列，性能远高于两层for循环实现。

import numpy as np

def greedy_row_match(arr):
    n_rows, n_cols = arr.shape
    # 预计算每行降序排列的列索引，numpy内置操作效率极高
    sorted_cols = np.argsort(-arr, axis=1)
    used_cols = np.zeros(n_cols, dtype=bool)
    res = []
    for row_idx in range(n_rows):
        # 遍历当前行排好序的列，找第一个未被使用的
        for col_idx in sorted_cols[row_idx]:
            if not used_cols[col_idx]:
                res.append((row_idx, col_idx))
                used_cols[col_idx] = True
                break
        # 所有列都用完了提前终止
        if len(res) == n_cols:
            break
    return res

示例测试

arr = np.array(
   [[0.86, 0.23, 0.83, 0.79],
   [0.15, 0.98, 0.86, 0.47],
   [1.  , 0.08, 0.01, 0.04],
   [0.78, 0.82, 0.17, 0.56],
   [0.73, 0.91, 0.52, 0.31]])
print(greedy_row_match(arr))
# 输出：[(0, 0), (1, 1), (2, 3), (3, 2)]，和示例结果完全一致

性能优化

如果需要进一步提升匹配速度，可以引入numba的JIT编译，整体速度可提升10倍以上：

from numba import jit

@jit(nopython=True)
def greedy_row_match_numba(arr, sorted_cols):
    n_rows, n_cols = arr.shape
    used_cols = np.zeros(n_cols, dtype=np.bool_)
    res = []
    for row_idx in range(n_rows):
        for col_idx in sorted_cols[row_idx]:
            if not used_cols[col_idx]:
                res.append((row_idx, col_idx))
                used_cols[col_idx] = True
                break
        if len(res) == n_cols:
            break
    return res

# 调用方式
sorted_cols = np.argsort(-arr, axis=1)
res = greedy_row_match_numba(arr, sorted_cols)

10000×10000的float矩阵下，np.argsort步骤约耗时1~2秒，后续匹配过程加numba后仅需几十毫秒，完全满足性能要求。

方案2：全局最大权匹配（可选） #

如果需要所有匹配对的数值总和全局最大（不按行顺序优先），可以使用scipy的线性指派求解器，该方案时间复杂度为O(n³)，仅适合规模小于2000×2000的矩阵：

from scipy.optimize import linear_sum_assignment

row_ind, col_ind = linear_sum_assignment(-arr) # 传入负矩阵求解最大值匹配
res = list(zip(row_ind, col_ind))

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