CODEWITHSUNDEEP
pythonnumpy
krasnaya
krasnaya

Reputation: 3135

Get indices of N maximum values in a numpy array, with random tie-breaking

I'm trying to get the top N maximum values in a numpy array, with random tie-breaking if the values are equal.

I can get the top N maximum values as follows (taken from here), but this code always returns the first "4" (ie, index 1). Is there a way to make it choose randomly amongst the 4s?

>>> a
array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> ind = np.argpartition(a, -4)[-4:]
>>> ind
array([1, 5, 8, 0])

Upvotes: 2

Views: 547

Answers (1)

Andy Hayden
Andy Hayden

Reputation: 375745

You could randomize the order before sorting, and reapply that same permutation:

In [11]: p = np.random.permutation(len(a))

In [12]: p[np.argpartition(a[p], -4)[-4:]]
Out[12]: array([7, 8, 0, 5])

Note: if we run this again we may get another solution:

In [13]: p = np.random.permutation(len(a))

In [14]: p[np.argpartition(a[p], -4)[-4:]]
Out[14]: array([1, 8, 0, 5])

Upvotes: 4

Related Questions