CODEWITHSUNDEEP
pythonnumpysortingmultidimensional-arraydistance
NeStack
NeStack

Reputation: 2014

Sort an array of multi D points by distance to a reference point

I have a reference point p_ref stored in a numpy array with a shape of (1024,), something like:

print(p_ref)
>>> array([ p1,  p2,  p3, ..., p_n])

I also have a numpy array A_points with a shape of (1024,5000) containing 5000 points, each having 1024 dimensions like p_ref. My problem: I would like to sort the points in A_points by their (eucledian) distance to p_ref!

How can I do this? I read about scipy.spatial.distance.cdist and scipy.spatial.KDTree, but they both weren't doing exactly what I wanted and when I tried to combine them I made a mess. Thanks!

For reference and consistency let's assume:

p_ref = np.array([0,1,2,3])
A_points = np.reshape(np.array([10,3,2,13,4,5,16,3,8,19,4,11]), (4,3))

Expected output:

array([[ 3,  2, 10],
       [ 4,  5, 13],
       [ 3,  8, 16],
       [ 4, 11, 19]])

Upvotes: 0

Views: 840

Answers (2)

Divakar
Divakar

Reputation: 221574

You can do something like this -

A_points[:,np.linalg.norm(A_points-p_ref[:,None],axis=0).argsort()]

Another with np.einsum that should be more efficient than np.linalg.norm -

d = A_points-p_ref[:,None]
out = A_points[:,np.einsum('ij,ij->j',d,d).argsort()]

Further optimize to leverage fast matrix-multiplication to replace last step -

A_points[:,((A_points**2).sum(0)+(p_ref**2).sum()-2*p_ref.dot(A_points)).argsort()]

Upvotes: 1

Christian Sloper
Christian Sloper

Reputation: 7510

EDIT: Updated on suggestions by the OP.

I hope I understand you correctly, but you can calculate the distance between two vectors by using numpy.linalg.norm. Using this it should be as simple as:

A_sorted = sorted( A_points.T, key = lambda x: np.linalg.norm(x - p_ref ) )
A_sorted = np.reshape(A_sorted, (3,4)).T

Upvotes: 2

Related Questions