How to sort a 2d numpy array based on their value when put into a function in python

Let’s say I have a NumPy array:

[[7 2]
 [7 3]
 [2 8]
 [4 3]
 [5 5]] 

Where the 0th index is the x value and the 1st index is the y value. How do I sort these values so that when I put them into the function:
(x^2 + y- 11)^2 + (x + y^2 -7)^2, they get sorted in ascending order depending on the results? so the sorted values would look like this:

[[4 3]
 [5 5]
 [7 2]
 [7 3]
 [2 8]]

The arrays can have duplicates.

One of my ideas would be to use the .argsort() method, though I don’t know how I could implement that.

Thanks!

>Solution :

You can apply the function you have along the first axis to get a one dimensional array with the function values. Passing that result to np.argsort() will give you the proper sorting indices:

a = np.array([
     [7, 2],
     [7, 3],
     [2, 8],
     [4, 3],
     [5, 5]] 
)

def my_func(row):
    x, y = row
    return (x ** 2 + y - 11) ** 2 + (x + y ** 2) ** 2

f = np.apply_along_axis(my_func, 1, a)
# array([1721, 1937, 4357,  233, 1261])

indices = np.argsort(f)
# array([3, 4, 0, 1, 2])

a[indices]
# array([[4, 3],
#        [5, 5],
#        [7, 2],
#        [7, 3],
#        [2, 8]])

Per @mozway’s comment…this is significanlty faster since it allows Numpy to vectorize the function:

x,y = a.T
aa = (x ** 2 + y - 11) ** 2 + (x + y ** 2) ** 2 
indices = np.argsort(aa)
a[indices]

with the same result.