Python Numpy中返回下标操作函数-节约时间的利器

如果觉得Python慢，那么首先应该想到是不是没有用对。

Numpy是Python中自带的一个数值计算库，包含了大量数值计算的常用方法。其底层大量使用C/C++(超过50%的代码量)，矩阵计算调用LAPACK库(Fortran)，同时在大量代码优化的层面做了工作，使得其内置方法速度奇快。

Numpy的设计也颇具人性化， 有许多具有特色又好用的方法可供使用。其中以arg开头函数就是专门为返回下标（集）而设计的，总共有5个这样的函数，它们简单易用，在许多问题里可以将程序极度简化，从而提高我们的工作效率。

1. argwhere()

这个函数与where基本一致，主要用于查找元素所在的位置：

a = np.array([46, 57, 23, 39, 1, 10, 0, 120])

# 找出a中大于1的所有元素的下标 
np.argwhere(a > 1)
Out: 
array([[0],
       [1],
       [2],
       [3],
       [5],
       [7]], dtype=int64)
# 找出a中所有小于10的元素的下标
np.argwhere(a < 10)
Out: 
array([[4],
       [6]], dtype=int64)

如果拿到一个需要返回元素下标的任务，我们从C直接考虑的话可能想到的是去遍历元素，再利用一个临时变量来计数，在匹配过程中返回这个计数值。虽然实现起来也很简单，但速度不是很快。而用这种方法，简单易懂，同时又超快。

2. argmin(),argmax()

最大，最小值的下标。这种操作应该也是我们非常常用的，而这两个函数的用法实际上和np.max(),np.min()的用法是完全一样的。

a = np.array([46, 57, 23, 39, 1, 10, 0, 120])
np.argmin(a)
Out: 6
np.argmax(a)
Out: 7

3. argsort()

这个函数非常厉害，它首先将输入的数组进行升序排列，再将排列后各元素的原下标依次返回。这种操作在一些需要联动排序的时候特别管用。比如我们之前讲过的在PCA的计算过程中，需要对特征值从大到小排序，同时要对它们对应的特征向量也排序。那么这个时候直接返回特征值排序后的原下标就能极大程序上简化我们的代码，同时又能节约大量的时间。

a = np.array([46, 57, 23, 39, 1, 10, 0, 120])

np.argsort(a)
Out: array([6, 4, 5, 2, 3, 0, 1, 7], dtype=int64)
    
a[np.argsort(a)]
Out: array([  0,   1,  10,  23,  39,  46,  57, 120])

上面最后一步主要是为了用它检验函数的正确性，同时也方便对照分析。

根据第4行的结果，我们来看前2个元素。它返回值表示：排序后的数组中第一个元素应该是原数组中的第7个元素0，第二个元素应该是原数组中的第5个元素1。那么看看第7行，完全正确。

当然，还有一个问题是np.sort本身只进行升序排列，那么当我们需要降序时直接将下标进行反序即可：

np.argsort(a)[::-1]
Out: array([7, 1, 0, 3, 2, 5, 4, 6], dtype=int64)

a[np.argsort(a)[::-1]]
Out: array([120,  57,  46,  39,  23,  10,   1,   0])

4. argpartition()

这也是一个非常神奇的函数，它主要是对应partition函数。如果只看文档说明很有可能理解错误。所以我们先简单说明一下函数的主要参数的含义。

np.argpartition(a,k)

这里a是 我们的原数组，k比如好理解的一种方式是：若将a进行排序，那么排序后的数组中的下标为k个元素应该处在其最终的位置。这句话有一点不太好理解。当然如果我们学过归并排序（快排的基础）的话，就可以理解这种分组的思路。

这个算法最终返回的下标里面，下标为k的元素一定处在原数组升序排列后的最终位置，它前面所有值全部小于等于它，后面所有值全部大于等于它，但其前后所有元素内的顺序是不关心的。所以这样就自然而然地能想到，由于它只考虑该元素的最终位置，那么时间复杂度可以在很大程度上降低。

先看一个简单的例子：

a = np.array([46, 57, 23, 39, 1, 10, 0, 120])

np.argpartition(a,3)
Out: array([6, 4, 5, 2, 3, 0, 1, 7], dtype=int64)
    
a[np.argpartition(a,3)]
Out: array([  0,   1,  10,  23,  39,  46,  57, 120])

这里可以清楚地看到，排序后的元素中下标为3的元素23的前面已经全部小于它，后面已经全部大于它，因此这个元素已经处在其排序后的最终位置了。

那么利用这个功能，我们主要是可以用它来找出元素中第k小的元素或者是第k大的元素所在的下标。

a[np.argpartition(a,3)[3]]
Out: 23		# 第四小的元素刚好是23
    
a[np.argpartition(a,-3)[-3]]
Out[]: 46	# 这里是倒数第3小的元素，也就是第3大的元素，正是46

另外，上述方法也可以指定具体的排序方法。比如上面提到的这种思想的基本排序方法是归并排序，其实也可以指定快速排序等方法。不过在数据量级不是特别大的情况下差异并不大。