2

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])
Community
  • 1
  • 1
krasnaya
  • 2,995
  • 3
  • 21
  • 19
  • It doesn't seem to be in any obvious order to begin with though? `np.argpartition([4, 4, 1, 4, 4], -2)[-2:]` gives `[1, 4]` while `np.argpartition([4, 1, 1, 4, 4], -2)[-2:]` gives `[3, 4]`. "in partitioned order" – endolith Aug 17 '22 at 21:20

1 Answers1

4

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])
Andy Hayden
  • 359,921
  • 101
  • 625
  • 535