1. 程式人生 > >Python庫Numpy的argpartition函式淺析

Python庫Numpy的argpartition函式淺析

最近在研究用Python實現K鄰近演算法的KD樹,發現一個奇怪的現象,就是為什麼沒有在網上搜到關於numpy的argpartition的任何中文的文章。難道是我搜索方法有問題嗎?所以自己寫一篇簡單的文章說說這個有意思的函式。
這個函式有什麼用呢?我舉幾個例子,您可能就明白她的作用了。
例子一:找到陣列的最大最小值,雖然這不是很合適,牛刀小用了。

>>> import numpy as np
>>> arr = np.array([46, 57, 23, 39, 1, 10, 0, 120])
>>> np.argpartition(arr, len(arr)-1
) array([5, 3, 2, 6, 4, 0, 1, 7], dtype=int32) >>> arr[np.argpartition(arr, len(arr)-1)[len(arr)-1]] 120 >>> arr[np.argpartition(arr, 0)[0]] 0

例子二:找到陣列的第3小(index=2)的值和第2大(index=-2)的值。這次就很合適了。

>>> arr[np.argpartition(arr, 2)[2]]
10
>>> arr[np.argpartition(arr, -2)[-2]]
57

例子三:同時找到第3和第4小的值。注意這裡,用[2,3]同時將第3和第4小的排序好,然後可以分別通過下標[2]和[3]取得。

>>> arr[np.argpartition(arr, [2,3])[2]]
10
>>> arr[np.argpartition(arr, [2,3])[3]]
23

有朋友您可能會問了,我直接把陣列arr排序一次不就行了嗎?還需要argpartition做什麼?這是個好問題。其實我開始也有這樣的疑問,然後我做了一個實驗,立刻明白了這樣做的好處。

如果您也感興趣,不妨聽我再多說幾句。這個實驗是這樣的。

我生成了一個數組,為了達到更好的效果,這個陣列被設計成倒序,這樣sort會得到最壞的結果。

>>> arr = np.array([8,7,6,5,4,3,2,1])
>>> np.argpartition(arr, 0)
array([7, 1, 2, 3, 4, 5, 6, 0], dtype=int32)
>>> np.argpartition(arr, 1)
array([7, 6, 2, 3, 4, 5, 1, 0], dtype=int32)
>>> np.argpartition(arr, 2)
array([7, 6, 5, 3, 4, 2, 1, 0], dtype=int32)
>>> np.argpartition(arr, 3)
array([6, 7, 5, 4, 3, 1, 2, 0], dtype=int32)
>>> np.argpartition(arr, 4)
array([4, 7, 6, 5, 3, 1, 2, 0], dtype=int32)
>>> np.argpartition(arr, 5)
array([4, 7, 6, 5, 3, 2, 1, 0], dtype=int32)
>>> np.argpartition(arr, 6)
array([4, 7, 6, 5, 3, 2, 1, 0], dtype=int32)
>>> np.argpartition(arr, 7)
array([4, 7, 6, 5, 3, 2, 1, 0], dtype=int32)

第一次呼叫,給第二個引數傳了0,說明我需要返回最小值得索引index。得到的返回值是array([7, 1, 2, 3, 4, 5, 6, 0], dtype=int32),在這個返回的array中,我關心的是第0個值(7),它是原陣列arr的索引,arr[7]就是我要找的最小值。請注意返回值中的其他幾個索引值,和原陣列的索引比起來,他們基本上沒有什麼變化。接下來的幾次呼叫也是這種情況,其實這也就說明argpartition沒有對他不關心的資料做太大的改動或者操作。

argpartition只排序第2個引數指定的位置,其他的位置的值不保證排序正確。因為這樣只調用一次argpartition函式,也就是執行了一遍搜尋就找到了n個值的排位,所以argpartition的效率比較高(別激動,我猜的)。據說是O(n)。

下面是函式的英文幫助,我在KD樹的程式碼裡只用到了前2個引數,而且只對一維陣列排序,夠了。使用不同的引數可以對複雜資料處理,我沒有用到,也沒有去實驗,以後用到了我會補充的。
我寫的KD樹程式碼:http://blog.csdn.net/weixin_37722024/article/details/62226957

numpy.argpartition
numpy.argpartition(a, kth, axis=-1, kind=’introselect’, order=None)[source]
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.

New in version 1.8.0.

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 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. In other words, a[index_array] yields a partitioned a.
See also
partition
Describes partition algorithms used.
ndarray.partition
Inplace partition.
argsort
Full indirect sort
Notes

See partition for notes on the different selection algorithms.

Examples

One dimensional array:

>

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

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