0

Is there any way to use argpartition on a general n-dimensional array and have it return only the first 2 components?

Of course I can just call argpartition on the array and then discard everything after the second component, but I need it to be faster. Is there a way to get numpy to compute (and return) only the first two components? If I remember correctly finding the first two components should be O(n) + O(n) as opposed to O(n log n) of sorting the whole thing; and there is a space advantage in not having to return a big array just to discard most of it immediately after

Thank you for your help!

Ant
  • 5,151
  • 2
  • 26
  • 43
  • Possible [duplicate](https://stackoverflow.com/questions/6910641/how-to-get-indices-of-n-maximum-values-in-a-numpy-array). use argpartition(a,2). You have to point out the multidimentionnal approach. – B. M. Jul 02 '17 at 20:32
  • @B.M. To my understanding, argpartition(a,2) will return an array with the same number of elements, just the first 2 in the correct order. I want it to just return the first 2 in the correct order, taking advantage of the smaller amount of computations it needs to perform – Ant Jul 02 '17 at 23:24

0 Answers0