找到两个“连接”矩阵的最大最小值的最快方法

Question

我想最大化以下功能：

f(i, j, k) = min(A(i, j), B(j, k))

其中A 和B 是矩阵，i、j 和k 是范围高达矩阵各自维度的索引。我想找到(i, j, k) 以使f(i, j, k)最大化。我目前这样做如下：

import numpy as np
import itertools

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])

# All the different i,j,k
combinations = itertools.product(np.arange(shape_a[0]), np.arange(shape_a[1]), np.arange(shape_b[1]))
combinations = np.asarray(list(combinations))

A_vals = A[combinations[:, 0], combinations[:, 1]]
B_vals = B[combinations[:, 1], combinations[:, 2]]

f = np.min([A_vals, B_vals], axis=0)

best_indices = combinations[np.argmax(f)]
print(best_indices)

[ 49  14 136]

这比遍历所有(i, j, k) 更快，但是很多（并且大部分）时间都花在了构建A_vals 和B_vals 矩阵上。这是不幸的，因为它们包含许多重复值，因为相同的 i、j 和 k 多次出现。有没有办法做到这一点，(1) numpy 的矩阵计算速度可以保持，(2) 我不必构造内存密集型 A_vals 和 B_vals 数组。

在其他语言中，您也许可以构造矩阵，以便它们包含指向 A 和 B 的指针，但我不知道如何在 Python 中实现这一点。

Answer 1

您的值i,j,k 由集合{A,B} 中的最大值的索引确定。您可以简单地使用np.argmax()。

if np.max(A) < np.max(B):
    ind = np.unravel_index(np.argmax(A),A.shape)
else:
    ind = np.unravel_index(np.argmax(B),B.shape)

它只会返回两个值，i,j 如果max({A,B}) = max({A}) 或j,k 如果max({A,B}) = max({B})。但是，如果你得到i,j，那么k 可以是任何符合B 数组形状的值，所以随机选择其中一个值。

如果您还需要最大化其他值，那么：

if np.max(A) < np.max(B):
    ind = np.unravel_index(np.argmax(A),A.shape)
    ind = ind + (np.argmax(B[ind[1],:]),)
    
else:
    ind = np.unravel_index(np.argmax(B),B.shape)
    ind = (np.argmax(A[:,ind[0]]),) + ind

Answer 2

也许您可以根据 min 和 max 的实际作用重新评估您如何看待问题。假设您有以下具体示例：

>>> np.random.seed(1)
>>> print(A := np.random.randint(10, size=(4, 5)))
[[5 8 9 5 0]
 [0 1 7 6 9]
 [2 4 5 2 4]
 [2 4 7 7 9]]
>>> print(B := np.random.randint(10, size=(5, 3)))
[[1 7 0]
 [6 9 9]
 [7 6 9]
 [1 0 1]
 [8 8 3]]

您正在寻找A 和B 中的一对数字，使得A 中的列与B 中的行相同，并且您将获得最大的较小数字。

对于任何一组数字，当您取两个最大的数字时，会出现最大的成对最小值。因此，您正在寻找A 的每一列、B 的行中的最大值，这些对中的最小值，然后是最大值。下面是一个相对简单的解决方案：

candidate_i = A.argmax(axis=0)
candidate_k = B.argmax(axis=1)
j = np.minimum(A[candidate_i, np.arange(A.shape[1])], B[np.arange(B.shape[0]), candidate_k]).argmax()

i = candidate_i[j]
k = candidate_k[j]

确实，你看到了

>>> i, j, k
(0, 2, 2)
>>> A[i, j]
9
>>> B[j, k]
9

如果发生冲突，argmax 将始终选择第一个选项。