获取ndarray最大最小N个index

numpy模块中有几个常用的函数，用来获取最大最小值的index，或者获取排序后的index。

np.argmax

argument max，获取ndarray中最大值的index：

>>> import numpy as np
>>> a = np.array((3,1,2,4,6,1,8,5,9,2))
>>> a
array([3, 1, 2, 4, 6, 1, 8, 5, 9, 2])
>>> np.argmax(a)
8
>>> a[8]
9

如果输入是个matrix，np.argmax默认会将其flatten，然后返回index。除非指定axis参数，如下是官方help中的case：

>>> a = np.arange(6).reshape(2,3) + 10
    >>> a
    array([[10, 11, 12],
           [13, 14, 15]])
    >>> np.argmax(a)
    5
    >>> np.argmax(a, axis=0)
    array([1, 1, 1])
    >>> np.argmax(a, axis=1)
    array([2, 2])

可以这样来理解numpy模块中许多函数的axis参数，along axis！axis=0，即沿着row操作，沿着row操作，就是操作column。axis=1，即沿着column操作，含义就是操作row。

np.argmin

argument min，此函数与np.argmax对应。

np.argsort

argument sort，返回ndarray排序后的index：

>>> a
array([3, 1, 2, 4, 6, 1, 8, 5, 9, 2])
>>> b = np.argsort(a)
>>> b
array([1, 5, 2, 9, 0, 3, 7, 4, 6, 8])
>>> [a[i] for i in b]
[1, 1, 2, 2, 3, 4, 5, 6, 8, 9]

多维ndarray的情况：

>>> aa = np.array(((3,1),(2,5),(6,3)))
>>> aa
array([[3, 1],
       [2, 5],
       [6, 3]])
>>> np.argsort(aa, axis=0)
array([[1, 0],
       [0, 2],
       [2, 1]])
>>> np.argsort(aa, axis=1)
array([[1, 0],
       [0, 1],
       [1, 0]])

np.argsort函数返回与输入一样shape的ndarray，只不过里面存放的全都是index。默认的排序顺序是从小到达。

np.argpartition

partition是分割的意思，表示将ndarray中数据的某一部分partition出来，当然这某一部分数据是排序后的，默认使用introselect算法，从N个数中，找出最小的K个，这样的问题。

>>> a
array([3, 1, 2, 4, 6, 1, 8, 5, 9, 2])
>>> b = np.argpartition(a,3)
>>> [a[i] for i in b]  # a[b]
[1, 1, 2, 2, 3, 4, 8, 5, 9, 6]
>>> a[b[:3]]
array([1, 1, 2])

前面3个index是保证的，后面的不保证！

np.argpartition函数在很多时候，都会比使用np.argsort要快，后者是 fully sorted，我们不一定需要，前者是 partly，只需要找到前面的 K 个，计算就停下了。

上例是找出最小的K个，下面是找出最大的K个：

>>> a
array([3, 1, 2, 4, 6, 1, 8, 5, 9, 2])
>>> b = np.argpartition(a,-3)
>>> a[b[-3:]]
array([6, 8, 9])

partition之后，在任意一个partition内的顺序是undefined。

还可以给np.argpartition函数输入 sequence of int，没搞定啥意思，测试结果也没看懂。

-- EOF --

相关文章

• 用numpy计算均值，方差，标准差
• numpy创建array的方法汇总
• 用numpy计算Pearson相关系数
• 打印很大的numpy.array
• 如果np.test有失败...
• np.linalg.norm的使用
• np.array的stack操作
• 查看numpy的数据类型
• numpy.random的常用接口
• numpy.ndarray基本概念
• ndarray中的条件判断
• np.concatenate函数
• 用numpy计算协方差（covariance）
• numpy.array变形记
• numpy的基本操作
• 如何遍历numpy.array？