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

Time:12-30

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!

CodePudding user response:

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.

CodePudding user response:

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 !

CodePudding user response:

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]]
  • Related