Search code examples
pythonarrayssortingnumpy

How does np.partition() interpret the argument kth?


I am trying to figure out how np.partition function works. For example, consider

arr = np.array([5, 4, 1, 0, -1, -3, -4, 0])

If I call np.partition(arr, kth=2), I get

np.array([-4, -3, -1, 0, 1, 4, 5, 0])

I expect that, after partition, the array will split into elements less than one, one, and elements greater than one. But the second zero is placed on the last array position, which isn't its right place after partition.


Solution

  • The documentation says:

    Creates a copy of the array with its elements rearranged in such a way that the value of the element in kth position is in the position it would be in a sorted array. All elements smaller than the kth element are moved before this element and all equal or greater are moved behind it. The ordering of the elements in the two partitions is undefined.

    In the example you give, you have selected 2th element of the sorted list (starting from zero), which is -1, and it seems to be in the right position if the array was sorted.