查找两个关联矩阵min(A(i,j),B(j,k))最大值的最快Python实现方法

核心优化思路 #

我们可以完全不需要枚举所有(i,j,k)组合，基于以下逻辑就能得到最优解：

对于固定的中间索引j，要让min(A[i,j], B[j,k])取到最大值，最优选择必然是取A第j列的最大值和B第j行的最大值，两者的最小值就是当前j能贡献的最大f值。全局最优解一定是所有j对应的最大f值中最高的那个，对应索引组合就是答案。

实现代码 #

import numpy as np

shape_a = (100, 150)
shape_b = (shape_a[1], 200)

A = np.random.rand(shape_a[0], shape_a[1])
B = np.random.rand(shape_b[0], shape_b[1])

# 对每个j，计算A第j列的最大值和对应i索引
max_A_per_j = A.max(axis=0)
best_i_per_j = A.argmax(axis=0)

# 对每个j，计算B第j行的最大值和对应k索引
max_B_per_j = B.max(axis=1)
best_k_per_j = B.argmax(axis=1)

# 计算每个j对应的最大f值
f_per_j = np.minimum(max_A_per_j, max_B_per_j)

# 找到最优的j
best_j = f_per_j.argmax()

# 得到最优三元组
best_indices = (best_i_per_j[best_j], best_j, best_k_per_j[best_j])
print(best_indices)

方案优势 #

• 完全保留numpy矩阵运算的速度优势，时间复杂度仅为O(I*J + J*K)，比原方案的O(I*J*K)快数十到上百倍，维度越大性能差距越明显
• 无需构造百万级大小的组合数组、A_vals和B_vals数组，所有中间数组的大小仅等于中间维度J的大小，内存占用可以忽略不计

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