理解 np.argpartition 工作的问题答案

【问题标题】：Problem with understanding of work of np.argpartition理解 np.argpartition 工作的问题
【发布时间】：2021-12-04 20:56:47
【问题描述】：

我在执行 np.argpartition 时遇到问题我有 nd.array

example = np.array([[5,6,7,3,4],[1,2,3,7,5],[6,7,4,2,3],[1,2,3,5,9],[2,3,6,1,2,]])
out: [[5 6 7 3 4]
      [1 2 3 7 5]
      [6 7 4 2 3]
      [1 2 3 5 9]
      [2 3 6 1 2]]

我可以通过 np.argsort 获取排序数组的索引

print(np.argsort(example))
out:
     [[3 4 0 1 2]
      [0 1 2 4 3]
      [3 4 2 0 1]
      [0 1 2 3 4]
      [3 0 4 1 2]]

我想使用 np.argsort 来节省一些时间来执行，因为我在这个数组的每一行中只需要 3 个排序元素。我使用这段代码来做到这一点：

print(np.argpartition(example, 3, axis=1))
out: [[3 4 0 1 2]
      [1 0 2 4 3]
      [3 4 2 0 1]
      [1 0 2 3 4]
      [3 4 0 1 2]]

我希望每行的前三个索引将匹配排序数组中的索引，但事实并非如此ю 那不起作用。我不明白我做错了什么。

【问题讨论】：

答案在 argpartition 的文档中，所以我否决了你的问题
如果你比较sort和partition，你的例子会更清楚。 “阅读”arg... 的值很困难，即使是 argsort。

标签： python arrays numpy sorting

【解决方案1】：

np.argpartition(example, k, axis=1) 不返回前 k 个元素的排序数组。它只返回索引，使得只有第 (k+1) 个元素被排序。如果您在输出中看到，只有第 4 个元素与 argsort() 匹配

如果你想要前三个排序的元素，你必须给出一个 k 参数的列表

index_array = np.argpartition(example, [0,1,2], axis=1)
print(np.take_along_axis(example,index_array, axis=1)) ##this will give you first 3 sorted elements

【讨论】：