ENH: add quickselect algorithm and expose it via partition

A partition sorts the kth element into its sorted order and moves all
smaller elements before the kth element and all equal or greater
elements behind it.
The ordering of all elements in the partitions is undefined.

It is implemented via the introselection algorithm which has worst case
linear complexity compared to a full sort that has linearithmic
complexity.
The introselect algorithm uses a quickselect with median of three pivot
and falls back to a quickselect with median of median of five pivot if
no sufficient progress is made.

The pivots used during the search for the wanted kth element can
optionally be stored and reused for further partitionings of the array.
This is used by the python interface if an array of kth is provided to
the partitions function. This improves the performance of median and
which need to select two elements if the size of the array is even. A
percentile function interpolating between values also profits from this.

String selection is implemented in terms of quicksort which has the same
properties as a selection for now.

1 files changed, 69 insertions, 0 deletions

diff --git a/numpy/add_newdocs.py b/numpy/add_newdocs.py
index 39dd2205e..5bb800488 100644
--- a/numpy/add_newdocs.py
+++ b/numpy/add_newdocs.py
@@ -3045,6 +3045,23 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('argsort',
     """))
 
 
+add_newdoc('numpy.core.multiarray', 'ndarray', ('argpartition',
+    """
+    a.argpartition(kth, axis=-1, kind='quickselect', order=None)
+
+    Returns the indices that would partition this array.
+
+    Refer to `numpy.argpartition` for full documentation.
+
+    .. versionadded:: 1.8.0
+
+    See Also
+    --------
+    numpy.argpartition : equivalent function
+
+    """))
+
+
 add_newdoc('numpy.core.multiarray', 'ndarray', ('astype',
     """
     a.astype(dtype, order='K', casting='unsafe', subok=True, copy=True)
@@ -4201,6 +4218,58 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('sort',
     """))
 
 
+add_newdoc('numpy.core.multiarray', 'ndarray', ('partition',
+    """
+    a.partition(kth, axis=-1, kind='introselect', order=None)
+
+    Rearranges the elements in the array in such a way that 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.
+
+    .. versionadded:: 1.8.0
+
+    Parameters
+    ----------
+    kth : int or sequence of ints
+        Element index to partition by. The kth element value will be in its
+        final sorted position and all smaller elements will be moved before it
+        and all equal or greater elements behind it.
+        The order all elements in the partitions is undefined.
+        If provided with a sequence of kth it will partition all elements
+        indexed by kth of them into their sorted position at once.
+    axis : int, optional
+        Axis along which to sort. Default is -1, which means sort along the
+        last axis.
+    kind : {'introselect'}, optional
+        Selection algorithm. Default is 'introselect'.
+    order : list, optional
+        When `a` is an array with fields defined, this argument specifies
+        which fields to compare first, second, etc.  Not all fields need be
+        specified.
+
+    See Also
+    --------
+    numpy.partition : Return a parititioned copy of an array.
+    argpartition : Indirect partition.
+
+    Notes
+    -----
+    See ``np.partition`` for notes on the different algorithms.
+
+    Examples
+    --------
+    >>> a = np.array([3, 4, 2, 1])
+    >>> a.partition(a, 3)
+    >>> a
+    array([2, 1, 3, 4])
+
+    >>> a.partition((1, 3))
+    array([1, 2, 3, 4])
+    """))
+
+
 add_newdoc('numpy.core.multiarray', 'ndarray', ('squeeze',
     """
     a.squeeze(axis=None)