获取ndarray最大最小N个index

2020年11月12日 / 9次阅读 / Last Modified 2020年11月12日
NumPy

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 --

本文链接:https://www.pynote.net/archives/2774

留言区

《获取ndarray最大最小N个index》有1条留言

电子邮件地址不会被公开。 必填项已用*标注

  • 麦新杰

    好些np函数的axis都可以设置为2-int-tuple,不理解? [回复]


前一篇:
后一篇:

More


©Copyright 麦新杰 Since 2019 Python笔记

go to top