如何在将二维 numpy 数组放入 python 中的 function 时根据它们的值对它们进行排序

Question

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!

Answer 1

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.

Answer 2

So this works:

def f(x, y):
return (x**2 + y- 11)**2 + (x + y**2 -7)**2


def sortTuples(TupleList):
    output = [0, 0, 0, 0, 0]
    indexList = []
    for i in TupleList:
        x = i[0]
        y = i[1]
        indexList.append(f(x, y))
    indexList.sort()
    for i in TupleList:
        output[indexList.index(f(i[0], i[1]))] = i
    return output

Hope you find a nicer way to do this !

Answer 3

At least for small arrays, sorted is competitive to np.argsort , (especially if lists suffice for your task):

out = sorted(arr.tolist(), key=lambda x: (x[0]**2+x[1]-11)**2+(x[0]+x[1]**2-7)**2)

Output:

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