Reputation: 843
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.
Upvotes: 2
Views: 4268
Reputation: 221684
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.
Upvotes: 1