Numpy operation for euclidean distance between multidimensional arrays

Question

I have two numpy arrays. 'A' of size w,h,2 and 'B' with n,2. In other words, A is a 2-dimensional array of 2D vectors while B is a 1D array of 2D vectors. What i want as a result is an array of size w,h,n. The last dimension is an n-dimensional vector where each of the components is the euclidean distance between the corresponding vector from A (denoted by the first two dimensions w and h) and the nth vector of B.

I know that i can just loop through w, h and n in python manually and calculate the distance for each element, but i like to know if there is a smart way to do that with numpy operations to increase performance.

I found some similar questions but unfortunately all of those use input arrays of the same dimensionality.

Answer 1

Approach #1

You could reshape A to 2D, use Scipy's cdist that expects 2D arrays as inputs, get those euclidean distances and finally reshape back to 3D.

Thus, an implementation would be -

from scipy.spatial.distance import cdist

out = cdist(A.reshape(-1,2),B).reshape(w,h,-1)

Approach #2

Since, the axis of reduction is of length 2 only, we can just slice the input arrays to save memory on intermediate arrays, like so -

np.sqrt((A[...,0,None] - B[:,0])**2 + (A[...,1,None] - B[:,1])**2)

Explanation on A[...,0,None] and A[...,1,None] :

With that None we are just introducing a new axis at the end of sliced A. Well, let's take a small example -

In [54]: A = np.random.randint(0,9,(4,5,2))

In [55]: A[...,0].shape
Out[55]: (4, 5)

In [56]: A[...,0,None].shape
Out[56]: (4, 5, 1)

In [57]: B = np.random.randint(0,9,(3,2))

In [58]: B[:,0].shape
Out[58]: (3,)

So, we have :

A[...,0,None]  : 4 x 5 x 1 
B[:,0]         :         3

That is essentially :

A[...,0,None]  : 4 x 5 x 1 
B[:,0]         : 1 x 1 x 3

When the subtraction is performed, the singleton dims are broadcasted corresponding to the dimensions of the other participating arrays -

A[...,0,None] - B  : 4 x 5 x 3

We repeat this for the second index along the last axis. We add these two arrays after squaring and finally a square-root to get the final eucl. distances.