广播的最小/最大 arrays

Question

Is there a way to efficiently compare multiple arrays which are broadcast together? For example:

a = np.arange( 0,  9).reshape(3,3)
b = np.arange( 9, 18).reshape(3,3)
c = np.arange(18, 27).reshape(3,3)

If I were to broadcast these as follows:

abc = a[:,:,None,None,None,None] + b[None,None,:,:,None,None] + c[None,None,None,None,:,:]

Then each element of abc is equal to a_ij + b_kl + c_mn where ij , kl , and mn index the respective arrays. What I would like instead is to get min(a_ij, b_kl, c_mn) , or ideally, max(a_ij, b_kl, c_mn) - min(a_ij, b_kl, c_mn) . Is there an efficient way in which I can do this?

I could, of course, broadcast temporary arrays as:

Abc =     a[:,:,None,None,None,None] + 0 * b[None,None,:,:,None,None] + 0 * c[None,None,None,None,:,:]
aBc = 0 * a[:,:,None,None,None,None] +     b[None,None,:,:,None,None] + 0 * c[None,None,None,None,:,:]
abC = 0 * a[:,:,None,None,None,None] + 0 * b[None,None,:,:,None,None] +     c[None,None,None,None,:,:]

and then find the min/max from these arrays, however, these arrays can get quite large. It would be better if there were some way to do it in one step.

And as an additional note, these arrays are guaranteed to be broadcastable, but not necessarily have the same shape (for example, (1, 3) and (3, 3) ).

Answer 1

You can store an intermediate array (smaller than your final result anyway) by doing the operation on a and b :

temp = np.minimum.outer(a.ravel(), b.ravel())
res = np.minimum.outer(temp.ravel(), c.ravel())

and then repeat that same operation with c . minimum computes element-wise min of 2 arrays. Since it is a ufunc , you can use outer to apply that operation to all pairs of values for those 2 arrays.

You can reshape res as you prefer.

Edit # 1

Thanks to P. Panzer comment, you do not need to use 1D arrays with ufunc.outer , which results in even simpler code:

temp = np.minimum.outer(a, b)
res = np.minimum.outer(temp, c)