NumPy的argsort函数能处理并列吗？

Question

我有一个numpy数组：

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

我希望得到前三个项目。调用：

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

返回

array([2, 0, 4])

注意，值 foo[1] 和 foo[4] 相等，因此numpy.argsort() 通过返回在数组中最后出现的项目的索引来处理平局；即索引为 4。

对于我的应用程序，我不能让平局总是偏向数组的末尾，那么如何实现随机平局？也就是说，一半的时间我会得到array([2, 0, 4])，另一半时间我会得到array([2, 0, 1])。

- BoltzmannBrain

---

使用lexsort或为每个值添加随机值来解决问题，参见如何使argsort结果在相等的值之间随机？ - Eric Tsui

谢谢，最终我选择了 numpy.lexsort((numpy.random.random(foo.size), foo))[::-1][:3]。 - BoltzmannBrain

Answer 1

以下是一种方法：

使用 numpy.unique 来对数组进行排序并删除重复项。传递 return_inverse 参数来获取按原始数组值排序的数组中给出原始数组值的索引。然后，通过查找逆数组中值等于该项在唯一数组中的索引的索引，可以得到所有绑定项的索引。

例如：

foo = array([3, 1, 4, 0, 1, 0])
foo_unique, foo_inverse = unique(foo, return_inverse=True)

# Put largest items first
foo_unique = foo_unique[::-1]
foo_inverse = -foo_inverse + len(foo_unique) - 1

foo_top3 = foo_unique[:3]

# Get the indices into foo of the top item
first_indices = (foo_inverse == 0).nonzero()

# Choose one at random
first_random_idx = random.choice(first_indices)

second_indices = (foo_inverse == 1).nonzero()
second_random_idx = random.choice(second_indices)

# And so on...

numpy.unique 是使用 argsort 实现的，因此查看其实现可能会建议一种更简单的方法。

- codewarrior

---

其实，我为什么要提到 numpy.unique 呢？你可以使用 (foo == foo[foo.argsort()[::-1][0]]).nonzero() 来获取所有第一名并列的数据。 - codewarrior

是的，与如何使argsort结果在相等值之间随机的答案相比，我的答案确实很愚蠢。 - codewarrior

实际上，你的回答也是一个很好的尝试。关键是几乎相同，都利用了随机。 - Eric Tsui