Search code examples
pythonarrayssortingproximity

sort numpy array elements by the value of a condition on the elements


I need to sort a numpy array of points by increasing distance from another point.

import numpy as np

def dist(i,j,ip,jp): 
    return np.sqrt((i-ip)**2+(j-jp)**2)

arr = np.array([[0,0],[1,2],[4,1]])

What I would like to do is use function dist(1,1,ip,jp) between a fixed point [i,j]=[1,1] and each ordered pair [ip,jp] in arr to return arr with each element sorted from lowest to highest proximity to [i,j]. Anyone have a quick solution for this?

The output I want is new_arr = np.array([[1,2],[0,0],[4,1]])

I have some ideas but they're wildly inefficient seeming.

Thanks!


Solution

  • I realize now this is a popular question, so years later, here's my own answer which uses the extremely powerful functionality of scipy.spatial. Here, scipy.spatial.cdist is used to do the distance computations. This is lightning fast and pythonic, without any "convert to list and convert back" hackery:

    from scipy.spatial.distance import cdist
    import numpy as np
    
    # EXAMPLE DATA
    arr = 20*np.random.random(size=(5000000,2))-10 # testing data
    pt = np.array([[1,1]]) # the point to eval proximity to
    
    # THE SOLUTION
    out = arr[np.argsort(cdist(arr,pt).squeeze())]
    

    Here, cdist gets an array of distances, squeeze kills the extra dimension in this array, argsort orders the indices into the distances by the distances, and arr[...] sorts arr by these indices.