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

Question

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!

Answer 1

There seem to be two ways to do this:

1. Convert the whole numpy array into a Python list, and sort it using Python's sort method with a key function.

    l = list(arr)
    l.sort(key=lambda coord: dist(1, 1, coord[0], coord[1]))
    arr = np.array(l)
   

2. Create a second numpy array by mapping dist() over the original array, use .argsort() to get the sorted order, then apply that to the original array.

    arr2 = np.vectorize(lambda coord: dist(1, 1, coord[0], coord[1]))(arr)
    arr3 = np.argsort(arr2)
    arr = np.array(arr)[arr3]