如何只取argpartition输出的第一个元素 - numpy

Question

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

当然我可以在数组上调用argpartition，然后在第二个组件之后丢弃所有内容，但我需要它更快。有没有办法让numpy只计算（并返回）前两个组件？如果我没记错的话，前两个组件应该是O（n）+ O（n）而不是O（n log n）对整个事物进行排序;

之后，不必立即丢弃大部分阵列就可以获得空间优势

感谢您的帮助！