CODEWITHSUNDEEP
pythonarraysnumpy
whizvids
whizvids

Reputation: 248

Can numpy argsort return lower index for ties?

I have a numpy array:

foo = array([3, 1, 4, 0, 1, 0])

I want the top 3 items. Calling

foo.argsort()[::-1][:3]

returns 

array([2, 0, 4])

Notice values foo[1] and foo[4] are equal, so numpy.argsort() handles the tie by returning the index of the item which appears last in the array; i.e. index 4.

For my application I want the tie breaking to return the index of the item which appears first in the array (index 1 here). How do I implement this efficiently?

Upvotes: 7

Views: 2404

Answers (2)

Julien
Julien

Reputation: 15081

What about simply this?

(-foo).argsort(kind='mergesort')[:3]

Why this works:

Argsorting in descending order (not what np.argsort does) is the same as argsorting in ascending order (what np.argsort does) the opposite values. You then just need to pick the first 3 sorted indices. Now all you need is make sure that the sort is stable, meaning in case of ties, keep first index first. NOTE: I thought the default kind=quicksort was stable but from the doc it appears only kind=mergesort is guaranteed to be stable: (https://docs.scipy.org/doc/numpy/reference/generated/numpy.sort.html)

The various sorting algorithms are characterized by their average speed, worst case performance, work space size, and whether they are stable. A stable sort keeps items with the same key in the same relative order. The three available algorithms have the following properties:

kind speed worst case work space stable

‘quicksort’ 1 O(n^2) 0 no

‘mergesort’ 2 O(n*log(n)) ~n/2 yes

‘heapsort’ 3 O(n*log(n)) 0 no

Upvotes: 6

Praveen
Praveen

Reputation: 7222

This is an extremely hacky answer, but why don't you just argsort the array in reverse? That way argsort picks the last index (in reverse), which is the first index.

This translates to:

>>> foo = np.array([3, 1, 4, 0, 1, 0])
>>> foo.argsort()[::-1]
array([2, 0, 4, 1, 5, 3])
>>> foo.size - 1 - foo[::-1].argsort()[::-1]
array([2, 0, 1, 4, 3, 5])

Upvotes: 2

Related Questions