CODEWITHSUNDEEP
pythonpython-3.xnumpysortingargs
Sreeragh A R
Sreeragh A R

Reputation: 3021

User defined tie breaker for argsort() in numpy

I have two arrays v and c (can read as value and cost).

I need to perform argsort() on v such that if 2 elements in v are the same, then they need to be sorted according to their corresponding elements in c.

Example

v = [4,1,4,4] # Here 0th, 2nd and 3rd elemnt are equal
c = [5,0,30,10]
numpy.argsort(v) = [1,0,2,3] # equal values sorted by index

Required output

[1,0,3,2] # c[0] < c[3] < c[2]

How to achieve this in Python?

Upvotes: 3

Views: 865

Answers (1)

Dani Mesejo
Dani Mesejo

Reputation: 61910

The function argsort receives an order parameter, from the docs:

When a is an array with fields defined, this argument specifies which fields to compare first, second, etc.

So you could create a structured array from the two values, and the pass the fields in order:

import numpy as np

v = [4, 1, 4, 4]  
c = [5, 0, 30, 10]

s = np.array(list(zip(v, c)), dtype=[('value', 'i4'), ('cost', 'i4')])
result = np.argsort(s, order=['value', 'cost'])

print(result)

Output

[1 0 3 2]

Upvotes: 4

Related Questions