Python 如何仅获取argpartition输出的第一个元素-numpy

有没有办法在一般的n维数组上使用argpartition，让它只返回前2个组件

当然，我可以只调用数组上的argpartition，然后丢弃第二个组件之后的所有内容，但我需要更快。有没有办法让numpy只计算（并返回）前两个组件？如果我没记错的话，前两个分量应该是O（n）+O（n），而不是整个排序的O（n logn）；而且还有一个空间优势，即不必返回一个大数组，然后立即丢弃其中的大部分

---

谢谢你的帮助

可能。使用argpartition（a，2）。您必须指出多维方法。@B.M.据我所知，argpartition（a，2）将返回一个元素数相同的数组，只返回正确顺序的前2个。我希望它只以正确的顺序返回前2个，利用它需要执行的计算量较小的优势