"""Module containing non-deprecated functions borrowed from Numeric.

"""

import functools

import math

import types

import numpy as np

from numpy._utils import set_module

from . import _methods, multiarray as mu, numerictypes as nt, overrides, umath as um

from ._multiarray_umath import _array_converter

from .multiarray import asanyarray, asarray, concatenate

_dt_ = nt.sctype2char

# functions that are methods

__all__ = [

'all', 'amax', 'amin', 'any', 'argmax',

'argmin', 'argpartition', 'argsort', 'around', 'choose', 'clip',

'compress', 'cumprod', 'cumsum', 'cumulative_prod', 'cumulative_sum',

'diagonal', 'mean', 'max', 'min', 'matrix_transpose',

'ndim', 'nonzero', 'partition', 'prod', 'ptp', 'put',

'ravel', 'repeat', 'reshape', 'resize', 'round',

'searchsorted', 'shape', 'size', 'sort', 'squeeze',

'std', 'sum', 'swapaxes', 'take', 'trace', 'transpose', 'var',

]

_gentype = types.GeneratorType

# save away Python sum

_sum_ = sum

array_function_dispatch = functools.partial(

overrides.array_function_dispatch, module='numpy')

# functions that are now methods

def _wrapit(obj, method, *args, **kwds):

conv = _array_converter(obj)

# As this already tried the method, subok is maybe quite reasonable here

# but this follows what was done before. TODO: revisit this.

arr, = conv.as_arrays(subok=False)

result = getattr(arr, method)(*args, **kwds)

return conv.wrap(result, to_scalar=False)

def _wrapfunc(obj, method, *args, **kwds):

bound = getattr(obj, method, None)

if bound is None:

return _wrapit(obj, method, *args, **kwds)

try:

return bound(*args, **kwds)

except TypeError:

# A TypeError occurs if the object does have such a method in its

# class, but its signature is not identical to that of NumPy's. This

# situation has occurred in the case of a downstream library like

# 'pandas'.

#

# Call _wrapit from within the except clause to ensure a potential

# exception has a traceback chain.

return _wrapit(obj, method, *args, **kwds)

def _wrapreduction(obj, ufunc, method, axis, dtype, out, **kwargs):

passkwargs = {k: v for k, v in kwargs.items()

if v is not np._NoValue}

if type(obj) is not mu.ndarray:

try:

reduction = getattr(obj, method)

except AttributeError:

pass

else:

# This branch is needed for reductions like any which don't

# support a dtype.

if dtype is not None:

return reduction(axis=axis, dtype=dtype, out=out, **passkwargs)

else:

return reduction(axis=axis, out=out, **passkwargs)

return ufunc.reduce(obj, axis, dtype, out, **passkwargs)

def _wrapreduction_any_all(obj, ufunc, method, axis, out, **kwargs):

# Same as above function, but dtype is always bool (but never passed on)

passkwargs = {k: v for k, v in kwargs.items()

if v is not np._NoValue}

if type(obj) is not mu.ndarray:

try:

reduction = getattr(obj, method)

except AttributeError:

pass

else:

return reduction(axis=axis, out=out, **passkwargs)

return ufunc.reduce(obj, axis, bool, out, **passkwargs)

def _take_dispatcher(a, indices, axis=None, out=None, mode=None):

return (a, out)

@array_function_dispatch(_take_dispatcher)

def take(a, indices, axis=None, out=None, mode='raise'):

"""

Take elements from an array along an axis.

When axis is not None, this function does the same thing as "fancy"

indexing (indexing arrays using arrays); however, it can be easier to use

if you need elements along a given axis. A call such as

``np.take(arr, indices, axis=3)`` is equivalent to

``arr[:,:,:,indices,...]``.

Explained without fancy indexing, this is equivalent to the following use

of `ndindex`, which sets each of ``ii``, ``jj``, and ``kk`` to a tuple of

indices::

Ni, Nk = a.shape[:axis], a.shape[axis+1:]

Nj = indices.shape

for ii in ndindex(Ni):

for jj in ndindex(Nj):

for kk in ndindex(Nk):

out[ii + jj + kk] = a[ii + (indices[jj],) + kk]

Parameters

----------

a : array_like (Ni..., M, Nk...)

The source array.

indices : array_like (Nj...)

The indices of the values to extract.

Also allow scalars for indices.

axis : int, optional

The axis over which to select values. By default, the flattened

input array is used.

out : ndarray, optional (Ni..., Nj..., Nk...)

If provided, the result will be placed in this array. It should

be of the appropriate shape and dtype. Note that `out` is always

buffered if `mode='raise'`; use other modes for better performance.

mode : {'raise', 'wrap', 'clip'}, optional

Specifies how out-of-bounds indices will behave.

* 'raise' -- raise an error (default)

* 'wrap' -- wrap around

* 'clip' -- clip to the range

'clip' mode means that all indices that are too large are replaced

by the index that addresses the last element along that axis. Note

that this disables indexing with negative numbers.

Returns

-------

out : ndarray (Ni..., Nj..., Nk...)

The returned array has the same type as `a`.

See Also

--------

compress : Take elements using a boolean mask

ndarray.take : equivalent method

take_along_axis : Take elements by matching the array and the index arrays

Notes

-----

By eliminating the inner loop in the description above, and using `s_` to

build simple slice objects, `take` can be expressed in terms of applying

fancy indexing to each 1-d slice::

Ni, Nk = a.shape[:axis], a.shape[axis+1:]

for ii in ndindex(Ni):

for kk in ndindex(Nk):

out[ii + s_[...,] + kk] = a[ii + s_[:,] + kk][indices]

For this reason, it is equivalent to (but faster than) the following use

of `apply_along_axis`::

out = np.apply_along_axis(lambda a_1d: a_1d[indices], axis, a)

Examples

--------

>>> import numpy as np

>>> a = [4, 3, 5, 7, 6, 8]

>>> indices = [0, 1, 4]

>>> np.take(a, indices)

array([4, 3, 6])

In this example if `a` is an ndarray, "fancy" indexing can be used.

>>> a = np.array(a)

>>> a[indices]

array([4, 3, 6])

If `indices` is not one dimensional, the output also has these dimensions.

>>> np.take(a, [[0, 1], [2, 3]])

array([[4, 3],

[5, 7]])

"""

return _wrapfunc(a, 'take', indices, axis=axis, out=out, mode=mode)

def _reshape_dispatcher(a, /, shape, order=None, *, copy=None):

return (a,)

@array_function_dispatch(_reshape_dispatcher)

def reshape(a, /, shape, order='C', *, copy=None):

"""

Gives a new shape to an array without changing its data.

Parameters

----------

a : array_like

Array to be reshaped.

shape : int or tuple of ints

The new shape should be compatible with the original shape. If

an integer, then the result will be a 1-D array of that length.

One shape dimension can be -1. In this case, the value is

inferred from the length of the array and remaining dimensions.

order : {'C', 'F', 'A'}, optional

Read the elements of ``a`` using this index order, and place the

elements into the reshaped array using this index order. 'C'

means to read / write the elements using C-like index order,

with the last axis index changing fastest, back to the first

axis index changing slowest. 'F' means to read / write the

elements using Fortran-like index order, with the first index

changing fastest, and the last index changing slowest. Note that

the 'C' and 'F' options take no account of the memory layout of

the underlying array, and only refer to the order of indexing.

'A' means to read / write the elements in Fortran-like index

order if ``a`` is Fortran *contiguous* in memory, C-like order

otherwise.

copy : bool, optional

If ``True``, then the array data is copied. If ``None``, a copy will

only be made if it's required by ``order``. For ``False`` it raises

a ``ValueError`` if a copy cannot be avoided. Default: ``None``.

Returns

-------

reshaped_array : ndarray

This will be a new view object if possible; otherwise, it will

be a copy. Note there is no guarantee of the *memory layout* (C- or

Fortran- contiguous) of the returned array.

See Also

--------

ndarray.reshape : Equivalent method.

Notes

-----

It is not always possible to change the shape of an array without copying

the data.

The ``order`` keyword gives the index ordering both for *fetching*

the values from ``a``, and then *placing* the values into the output

array. For example, let's say you have an array:

>>> a = np.arange(6).reshape((3, 2))

>>> a

array([[0, 1],

[2, 3],

[4, 5]])

You can think of reshaping as first raveling the array (using the given

index order), then inserting the elements from the raveled array into the

new array using the same kind of index ordering as was used for the

raveling.

>>> np.reshape(a, (2, 3)) # C-like index ordering

array([[0, 1, 2],

[3, 4, 5]])

>>> np.reshape(np.ravel(a), (2, 3)) # equivalent to C ravel then C reshape

array([[0, 1, 2],

[3, 4, 5]])

>>> np.reshape(a, (2, 3), order='F') # Fortran-like index ordering

array([[0, 4, 3],

[2, 1, 5]])

>>> np.reshape(np.ravel(a, order='F'), (2, 3), order='F')

array([[0, 4, 3],

[2, 1, 5]])

Examples

--------

>>> import numpy as np

>>> a = np.array([[1,2,3], [4,5,6]])

>>> np.reshape(a, 6)

array([1, 2, 3, 4, 5, 6])

>>> np.reshape(a, 6, order='F')

array([1, 4, 2, 5, 3, 6])

>>> np.reshape(a, (3,-1)) # the unspecified value is inferred to be 2

array([[1, 2],

[3, 4],

[5, 6]])

"""

if copy is not None:

return _wrapfunc(a, 'reshape', shape, order=order, copy=copy)

return _wrapfunc(a, 'reshape', shape, order=order)

def _choose_dispatcher(a, choices, out=None, mode=None):

yield a

yield from choices

yield out

@array_function_dispatch(_choose_dispatcher)

def choose(a, choices, out=None, mode='raise'):

"""

Construct an array from an index array and a list of arrays to choose from.

First of all, if confused or uncertain, definitely look at the Examples -

in its full generality, this function is less simple than it might

seem from the following code description::

np.choose(a,c) == np.array([c[a[I]][I] for I in np.ndindex(a.shape)])

But this omits some subtleties. Here is a fully general summary:

Given an "index" array (`a`) of integers and a sequence of ``n`` arrays

(`choices`), `a` and each choice array are first broadcast, as necessary,

to arrays of a common shape; calling these *Ba* and *Bchoices[i], i =

0,...,n-1* we have that, necessarily, ``Ba.shape == Bchoices[i].shape``

for each ``i``. Then, a new array with shape ``Ba.shape`` is created as

follows:

* if ``mode='raise'`` (the default), then, first of all, each element of

``a`` (and thus ``Ba``) must be in the range ``[0, n-1]``; now, suppose

that ``i`` (in that range) is the value at the ``(j0, j1, ..., jm)``

position in ``Ba`` - then the value at the same position in the new array

is the value in ``Bchoices[i]`` at that same position;

* if ``mode='wrap'``, values in `a` (and thus `Ba`) may be any (signed)

integer; modular arithmetic is used to map integers outside the range

`[0, n-1]` back into that range; and then the new array is constructed

as above;

* if ``mode='clip'``, values in `a` (and thus ``Ba``) may be any (signed)

integer; negative integers are mapped to 0; values greater than ``n-1``

are mapped to ``n-1``; and then the new array is constructed as above.

Parameters

----------

a : int array

This array must contain integers in ``[0, n-1]``, where ``n`` is the

number of choices, unless ``mode=wrap`` or ``mode=clip``, in which

cases any integers are permissible.

choices : sequence of arrays

Choice arrays. `a` and all of the choices must be broadcastable to the

same shape. If `choices` is itself an array (not recommended), then

its outermost dimension (i.e., the one corresponding to

``choices.shape[0]``) is taken as defining the "sequence".

out : array, optional

If provided, the result will be inserted into this array. It should

be of the appropriate shape and dtype. Note that `out` is always

buffered if ``mode='raise'``; use other modes for better performance.

mode : {'raise' (default), 'wrap', 'clip'}, optional

Specifies how indices outside ``[0, n-1]`` will be treated:

* 'raise' : an exception is raised

* 'wrap' : value becomes value mod ``n``

* 'clip' : values < 0 are mapped to 0, values > n-1 are mapped to n-1

Returns

-------

merged_array : array

The merged result.

Raises

------

ValueError: shape mismatch

If `a` and each choice array are not all broadcastable to the same

shape.

See Also

--------

ndarray.choose : equivalent method

numpy.take_along_axis : Preferable if `choices` is an array

Notes

-----

To reduce the chance of misinterpretation, even though the following

"abuse" is nominally supported, `choices` should neither be, nor be

thought of as, a single array, i.e., the outermost sequence-like container

should be either a list or a tuple.

Examples

--------

>>> import numpy as np

>>> choices = [[0, 1, 2, 3], [10, 11, 12, 13],

... [20, 21, 22, 23], [30, 31, 32, 33]]

>>> np.choose([2, 3, 1, 0], choices

... # the first element of the result will be the first element of the

... # third (2+1) "array" in choices, namely, 20; the second element

... # will be the second element of the fourth (3+1) choice array, i.e.,

... # 31, etc.

... )

array([20, 31, 12, 3])

>>> np.choose([2, 4, 1, 0], choices, mode='clip') # 4 goes to 3 (4-1)

array([20, 31, 12, 3])

>>> # because there are 4 choice arrays

>>> np.choose([2, 4, 1, 0], choices, mode='wrap') # 4 goes to (4 mod 4)

array([20, 1, 12, 3])

>>> # i.e., 0

A couple examples illustrating how choose broadcasts:

>>> a = [[1, 0, 1], [0, 1, 0], [1, 0, 1]]

>>> choices = [-10, 10]

>>> np.choose(a, choices)

array([[ 10, -10, 10],

[-10, 10, -10],

[ 10, -10, 10]])

>>> # With thanks to Anne Archibald

>>> a = np.array([0, 1]).reshape((2,1,1))

>>> c1 = np.array([1, 2, 3]).reshape((1,3,1))

>>> c2 = np.array([-1, -2, -3, -4, -5]).reshape((1,1,5))

>>> np.choose(a, (c1, c2)) # result is 2x3x5, res[0,:,:]=c1, res[1,:,:]=c2

array([[[ 1, 1, 1, 1, 1],

[ 2, 2, 2, 2, 2],

[ 3, 3, 3, 3, 3]],

[[-1, -2, -3, -4, -5],

[-1, -2, -3, -4, -5],

[-1, -2, -3, -4, -5]]])

"""

return _wrapfunc(a, 'choose', choices, out=out, mode=mode)

def _repeat_dispatcher(a, repeats, axis=None):

return (a,)

@array_function_dispatch(_repeat_dispatcher)

def repeat(a, repeats, axis=None):

"""

Repeat each element of an array after themselves

Parameters

----------

a : array_like

Input array.

repeats : int or array of ints

The number of repetitions for each element. `repeats` is broadcasted

to fit the shape of the given axis.

axis : int, optional

The axis along which to repeat values. By default, use the

flattened input array, and return a flat output array.

Returns

-------

repeated_array : ndarray

Output array which has the same shape as `a`, except along

the given axis.

See Also

--------

tile : Tile an array.

unique : Find the unique elements of an array.

Examples

--------

>>> import numpy as np

>>> np.repeat(3, 4)

array([3, 3, 3, 3])

>>> x = np.array([[1,2],[3,4]])

>>> np.repeat(x, 2)

array([1, 1, 2, 2, 3, 3, 4, 4])

>>> np.repeat(x, 3, axis=1)

array([[1, 1, 1, 2, 2, 2],

[3, 3, 3, 4, 4, 4]])

>>> np.repeat(x, [1, 2], axis=0)

array([[1, 2],

[3, 4],

[3, 4]])

"""

return _wrapfunc(a, 'repeat', repeats, axis=axis)

def _put_dispatcher(a, ind, v, mode=None):

return (a, ind, v)

@array_function_dispatch(_put_dispatcher)

def put(a, ind, v, mode='raise'):

"""

Replaces specified elements of an array with given values.

The indexing works on the flattened target array. `put` is roughly

equivalent to:

::

a.flat[ind] = v

Parameters

----------

a : ndarray

Target array.

ind : array_like

Target indices, interpreted as integers.

v : array_like

Values to place in `a` at target indices. If `v` is shorter than

`ind` it will be repeated as necessary.

mode : {'raise', 'wrap', 'clip'}, optional

Specifies how out-of-bounds indices will behave.

* 'raise' -- raise an error (default)

* 'wrap' -- wrap around

* 'clip' -- clip to the range

'clip' mode means that all indices that are too large are replaced

by the index that addresses the last element along that axis. Note

that this disables indexing with negative numbers. In 'raise' mode,

if an exception occurs the target array may still be modified.

See Also

--------

putmask, place

put_along_axis : Put elements by matching the array and the index arrays

Examples

--------

>>> import numpy as np

>>> a = np.arange(5)

>>> np.put(a, [0, 2], [-44, -55])

>>> a

array([-44, 1, -55, 3, 4])

>>> a = np.arange(5)

>>> np.put(a, 22, -5, mode='clip')

>>> a

array([ 0, 1, 2, 3, -5])

"""

try:

put = a.put

except AttributeError as e:

raise TypeError(f"argument 1 must be numpy.ndarray, not {type(a)}") from e

return put(ind, v, mode=mode)

def _swapaxes_dispatcher(a, axis1, axis2):

return (a,)

@array_function_dispatch(_swapaxes_dispatcher)

def swapaxes(a, axis1, axis2):

"""

Interchange two axes of an array.

Parameters

----------

a : array_like

Input array.

axis1 : int

First axis.

axis2 : int

Second axis.

Returns

-------

a_swapped : ndarray

For NumPy >= 1.10.0, if `a` is an ndarray, then a view of `a` is

returned; otherwise a new array is created. For earlier NumPy

versions a view of `a` is returned only if the order of the

axes is changed, otherwise the input array is returned.

Examples

--------

>>> import numpy as np

>>> x = np.array([[1,2,3]])

>>> np.swapaxes(x,0,1)

array([[1],

[2],

[3]])

>>> x = np.array([[[0,1],[2,3]],[[4,5],[6,7]]])

>>> x

array([[[0, 1],

[2, 3]],

[[4, 5],

[6, 7]]])

>>> np.swapaxes(x,0,2)

array([[[0, 4],

[2, 6]],

[[1, 5],

[3, 7]]])

"""

return _wrapfunc(a, 'swapaxes', axis1, axis2)

def _transpose_dispatcher(a, axes=None):

return (a,)

@array_function_dispatch(_transpose_dispatcher)

def transpose(a, axes=None):

"""

Returns an array with axes transposed.

For a 1-D array, this returns an unchanged view of the original array, as a

transposed vector is simply the same vector.

To convert a 1-D array into a 2-D column vector, an additional dimension

must be added, e.g., ``np.atleast_2d(a).T`` achieves this, as does

``a[:, np.newaxis]``.

For a 2-D array, this is the standard matrix transpose.

For an n-D array, if axes are given, their order indicates how the

axes are permuted (see Examples). If axes are not provided, then

``transpose(a).shape == a.shape[::-1]``.

Parameters

----------

a : array_like

Input array.

axes : tuple or list of ints, optional

If specified, it must be a tuple or list which contains a permutation

of [0, 1, ..., N-1] where N is the number of axes of `a`. Negative

indices can also be used to specify axes. The i-th axis of the returned

array will correspond to the axis numbered ``axes[i]`` of the input.

If not specified, defaults to ``range(a.ndim)[::-1]``, which reverses

the order of the axes.

Returns

-------

p : ndarray

`a` with its axes permuted. A view is returned whenever possible.

See Also

--------

ndarray.transpose : Equivalent method.

moveaxis : Move axes of an array to new positions.

argsort : Return the indices that would sort an array.

Notes

-----

Use ``transpose(a, argsort(axes))`` to invert the transposition of tensors

when using the `axes` keyword argument.

Examples

--------

>>> import numpy as np

>>> a = np.array([[1, 2], [3, 4]])

>>> a

array([[1, 2],

[3, 4]])

>>> np.transpose(a)

array([[1, 3],

[2, 4]])

>>> a = np.array([1, 2, 3, 4])

>>> a

array([1, 2, 3, 4])

>>> np.transpose(a)

array([1, 2, 3, 4])

>>> a = np.ones((1, 2, 3))

>>> np.transpose(a, (1, 0, 2)).shape

(2, 1, 3)

>>> a = np.ones((2, 3, 4, 5))

>>> np.transpose(a).shape

(5, 4, 3, 2)

>>> a = np.arange(3*4*5).reshape((3, 4, 5))

>>> np.transpose(a, (-1, 0, -2)).shape

(5, 3, 4)

"""

return _wrapfunc(a, 'transpose', axes)

def _matrix_transpose_dispatcher(x):

return (x,)

@array_function_dispatch(_matrix_transpose_dispatcher)

def matrix_transpose(x, /):

"""

Transposes a matrix (or a stack of matrices) ``x``.

This function is Array API compatible.

Parameters

----------

x : array_like

Input array having shape (..., M, N) and whose two innermost

dimensions form ``MxN`` matrices.

Returns

-------

out : ndarray

An array containing the transpose for each matrix and having shape

(..., N, M).

See Also

--------

transpose : Generic transpose method.

Examples

--------

>>> import numpy as np

>>> np.matrix_transpose([[1, 2], [3, 4]])

array([[1, 3],

[2, 4]])

>>> np.matrix_transpose([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])

array([[[1, 3],

[2, 4]],

[[5, 7],

[6, 8]]])

"""

x = asanyarray(x)

if x.ndim < 2:

raise ValueError(

f"Input array must be at least 2-dimensional, but it is {x.ndim}"

)

return swapaxes(x, -1, -2)

def _partition_dispatcher(a, kth, axis=None, kind=None, order=None):

return (a,)

@array_function_dispatch(_partition_dispatcher)

def partition(a, kth, axis=-1, kind='introselect', order=None):

"""

Return a partitioned copy of an array.

Creates a copy of the array and partially sorts it in such a way that

the value of the element in k-th position is in the position it would be

in a sorted array. In the output array, all elements smaller than the k-th

element are located to the left of this element and all equal or greater

are located to its right. The ordering of the elements in the two

partitions on the either side of the k-th element in the output array is

undefined.

Parameters

----------

a : array_like

Array to be sorted.

kth : int or sequence of ints

Element index to partition by. The k-th value of the element

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 of all elements in the partitions is undefined. If

provided with a sequence of k-th it will partition all elements

indexed by k-th of them into their sorted position at once.

axis : int or None, optional

Axis along which to sort. If None, the array is flattened before

sorting. The default is -1, which sorts along the last axis.

kind : {'introselect'}, optional

Selection algorithm. Default is 'introselect'.

order : str or list of str, optional

When `a` is an array with fields defined, this argument

specifies which fields to compare first, second, etc. A single

field can be specified as a string. Not all fields need be

specified, but unspecified fields will still be used, in the

order in which they come up in the dtype, to break ties.

Returns

-------

partitioned_array : ndarray

Array of the same type and shape as `a`.

See Also

--------

ndarray.partition : Method to sort an array in-place.

argpartition : Indirect partition.

sort : Full sorting

Notes

-----

The various selection algorithms are characterized by their average

speed, worst case performance, work space size, and whether they are

stable. A stable sort keeps items with the same key in the same

relative order. The available algorithms have the following

properties:

================= ======= ============= ============ =======

kind speed worst case work space stable

================= ======= ============= ============ =======

'introselect' 1 O(n) 0 no

================= ======= ============= ============ =======

All the partition algorithms make temporary copies of the data when

partitioning along any but the last axis. Consequently,

partitioning along the last axis is faster and uses less space than

partitioning along any other axis.

The sort order for complex numbers is lexicographic. If both the

real and imaginary parts are non-nan then the order is determined by

the real parts except when they are equal, in which case the order

is determined by the imaginary parts.

The sort order of ``np.nan`` is bigger than ``np.inf``.

Examples

--------

>>> import numpy as np

>>> a = np.array([7, 1, 7, 7, 1, 5, 7, 2, 3, 2, 6, 2, 3, 0])

>>> p = np.partition(a, 4)

>>> p

array([0, 1, 2, 1, 2, 5, 2, 3, 3, 6, 7, 7, 7, 7]) # may vary

``p[4]`` is 2; all elements in ``p[:4]`` are less than or equal

to ``p[4]``, and all elements in ``p[5:]`` are greater than or

equal to ``p[4]``. The partition is::

[0, 1, 2, 1], [2], [5, 2, 3, 3, 6, 7, 7, 7, 7]

The next example shows the use of multiple values passed to `kth`.

>>> p2 = np.partition(a, (4, 8))

>>> p2

array([0, 1, 2, 1, 2, 3, 3, 2, 5, 6, 7, 7, 7, 7])

``p2[4]`` is 2 and ``p2[8]`` is 5. All elements in ``p2[:4]``

are less than or equal to ``p2[4]``, all elements in ``p2[5:8]``

are greater than or equal to ``p2[4]`` and less than or equal to

``p2[8]``, and all elements in ``p2[9:]`` are greater than or

equal to ``p2[8]``. The partition is::

[0, 1, 2, 1], [2], [3, 3, 2], [5], [6, 7, 7, 7, 7]

"""

if axis is None:

# flatten returns (1, N) for np.matrix, so always use the last axis

a = asanyarray(a).flatten()

axis = -1

else:

a = asanyarray(a).copy(order="K")

a.partition(kth, axis=axis, kind=kind, order=order)

return a

def _argpartition_dispatcher(a, kth, axis=None, kind=None, order=None):

return (a,)

@array_function_dispatch(_argpartition_dispatcher)

def argpartition(a, kth, axis=-1, kind='introselect', order=None):

"""

Perform an indirect partition along the given axis using the

algorithm specified by the `kind` keyword. It returns an array of

indices of the same shape as `a` that index data along the given

axis in partitioned order.

Parameters

----------

a : array_like

Array to sort.

kth : int or sequence of ints

Element index to partition by. The k-th element will be in its

final sorted position and all smaller elements will be moved

before it and all larger elements behind it. The order of all

elements in the partitions is undefined. If provided with a

sequence of k-th it will partition all of them into their sorted

position at once.

axis : int or None, optional

Axis along which to sort. The default is -1 (the last axis). If

None, the flattened array is used.

kind : {'introselect'}, optional

Selection algorithm. Default is 'introselect'

order : str or list of str, optional

When `a` is an array with fields defined, this argument

specifies which fields to compare first, second, etc. A single

field can be specified as a string, and not all fields need be

specified, but unspecified fields will still be used, in the

order in which they come up in the dtype, to break ties.

Returns

-------

index_array : ndarray, int

Array of indices that partition `a` along the specified axis.

If `a` is one-dimensional, ``a[index_array]`` yields a partitioned `a`.

More generally, ``np.take_along_axis(a, index_array, axis=axis)``

always yields the partitioned `a`, irrespective of dimensionality.

See Also

--------

partition : Describes partition algorithms used.

ndarray.partition : Inplace partition.

argsort : Full indirect sort.

take_along_axis : Apply ``index_array`` from argpartition

to an array as if by calling partition.

Notes

-----

The returned indices are not guaranteed to be sorted according to

the values. Furthermore, the default selection algorithm ``introselect``

is unstable, and hence the returned indices are not guaranteed

to be the earliest/latest occurrence of the element.

`argpartition` works for real/complex inputs with nan values,

see `partition` for notes on the enhanced sort order and

different selection algorithms.

Examples

--------

One dimensional array:

>>> import numpy as np

>>> x = np.array([3, 4, 2, 1])

>>> x[np.argpartition(x, 3)]

array([2, 1, 3, 4]) # may vary

>>> x[np.argpartition(x, (1, 3))]

array([1, 2, 3, 4]) # may vary

>>> x = [3, 4, 2, 1]

>>> np.array(x)[np.argpartition(x, 3)]

array([2, 1, 3, 4]) # may vary

Multi-dimensional array:

>>> x = np.array([[3, 4, 2], [1, 3, 1]])

>>> index_array = np.argpartition(x, kth=1, axis=-1)

>>> # below is the same as np.partition(x, kth=1)

>>> np.take_along_axis(x, index_array, axis=-1)

array([[2, 3, 4],

[1, 1, 3]])

"""

return _wrapfunc(a, 'argpartition', kth, axis=axis, kind=kind, order=order)

def _sort_dispatcher(a, axis=None, kind=None, order=None, *, stable=None):

return (a,)

@array_function_dispatch(_sort_dispatcher)

def sort(a, axis=-1, kind=None, order=None, *, stable=None):

"""

Return a sorted copy of an array.

Parameters

----------

a : array_like

Array to be sorted.

axis : int or None, optional

Axis along which to sort. If None, the array is flattened before

sorting. The default is -1, which sorts along the last axis.

kind : {'quicksort', 'mergesort', 'heapsort', 'stable'}, optional

Sorting algorithm. The default is 'quicksort'. Note that both 'stable'

and 'mergesort' use timsort or radix sort under the covers and,

in general, the actual implementation will vary with data type.

The 'mergesort' option is retained for backwards compatibility.

order : str or list of str, optional

When `a` is an array with fields defined, this argument specifies

which fields to compare first, second, etc. A single field can

be specified as a string, and not all fields need be specified,

but unspecified fields will still be used, in the order in which

they come up in the dtype, to break ties.

stable : bool, optional

Sort stability. If ``True``, the returned array will maintain

the relative order of ``a`` values which compare as equal.

If ``False`` or ``None``, this is not guaranteed. Internally,

this option selects ``kind='stable'``. Default: ``None``.

.. versionadded:: 2.0.0

Returns

-------

sorted_array : ndarray

Array of the same type and shape as `a`.

See Also

--------

ndarray.sort : Method to sort an array in-place.

argsort : Indirect sort.

lexsort : Indirect stable sort on multiple keys.

searchsorted : Find elements in a sorted array.

partition : Partial sort.

Notes

-----

The various sorting algorithms are characterized by their average speed,

worst case performance, work space size, and whether they are stable. A

stable sort keeps items with the same key in the same relative

order. The four algorithms implemented in NumPy have the following

properties:

=========== ======= ============= ============ ========

kind speed worst case work space stable

=========== ======= ============= ============ ========

'quicksort' 1 O(n^2) 0 no

'heapsort' 3 O(n*log(n)) 0 no

'mergesort' 2 O(n*log(n)) ~n/2 yes

'timsort' 2 O(n*log(n)) ~n/2 yes

=========== ======= ============= ============ ========

.. note:: The datatype determines which of 'mergesort' or 'timsort'