[FIXED] Fastest way to find the maximum minimum value of two 'connected' matrices

Issue

I want to maximize the following function:

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

Where A and B are matrices and i, j and k are indices that range up to the respective dimensions of the matrices. I would like to find (i, j, k) such that f(i, j, k) is maximized. I am currently doing that as follows:

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]

This is faster than iterating over all (i, j, k), but a lot of (and most of the) time is spent constructing the A_vals and B_vals matrices. This is unfortunate, because they contain many many duplicate values as the same i, j and k appear multiple times. Is there a way to do this where (1) the speed of numpy's matrix computation can be preserved and (2) I don't have to construct the memory-intensive A_vals and B_vals arrays.

In other languages you could maybe construct the matrices so that they container pointers to A and B, but I do not see how to achieve this in Python.

Solution

Perhaps you could re-evaluate how you look at the problem in context of what min and max actually do. Say you have the following concrete example:

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

You are looking for a pair of numbers in A and B such that the column in A is the same as the row of B, and the you get the maximum smaller number.

For any set of numbers, the largest pairwise minimum happens when you take the two largest numbers. You are therefore looking for the max in each column of A, row of B, the minimum of those pairs, and then the maximum of that. Here is a relatively simple formulation of the solution:

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]

And indeed, you see that

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

If there are collisions, argmax will always pick the first option.

Answered By - Mad Physicist