np.argpartition 難以解決topK
topK是常用的一個功能，在python中，numpy等計算庫使用了豐富的底層優化，對于矩陣計算的效率遠高于python的for-loop實現。因此，我們希望盡量用一些numpy函數的組合實現topK。
pytorch 庫提供了topk函數，可以將高維數組沿某一維度（該維度共N項），選出最大（最小）的K項并排序。返回排序結果和index信息。奇怪的是，更輕量級的numpy庫并沒有直接提供 topK 函數。numpy只提供了argpartition 和 partition，可以將最大（最小）的K項排到前K位。以argpartition為例，最小的3項排到了前3位：
>>> x = np.array([3, 5, 6, 4, 2, 7, 1])
>>> x[np.argpartition(x, 3)]
array([2, 1, 3, 4, 5, 7, 6])
注意，argpartition實現的是 partial sorting，如上例，前3項和其余項被分開，但是兩部分各自都是不排序的！而我們可能更想要topK的幾項排好序（其余項則不作要求）。因此，下面提供一種基于argpartition的topK方法。
一個naive方法
最簡單的方法自然是全排序，然后取前K項。缺點在于，要把topK之外的數據也進行排序，當K N時較為浪費時間，復雜度為O ( n log n ) O(n \log n)O(nlogn):
def naive_arg_topK(matrix, K, axis=0):
"""
perform topK based on np.argsort
:param matrix: to be sorted
:param K: select and sort the top K items
:param axis: dimension to be sorted.
:return:
"""
full_sort = np.argsort(matrix, axis=axis)
return full_sort.take(np.arange(K), axis=axis)
# Example
>>> dists = np.random.permutation(np.arange(30)).reshape(6, 5)
array([[17, 28, 1, 24, 23, 8],
[ 9, 21, 3, 22, 4, 5],
[19, 12, 26, 11, 13, 27],
[10, 15, 18, 14, 7, 16],
[ 0, 25, 29, 2, 6, 20]])
>>> naive_arg_topK(dists, 2, axis=0)
array([[4, 2, 0, 4, 1, 1],
[1, 3, 1, 2, 4, 0]])
>>> naive_arg_topK(dists, 2, axis=1)
array([[2, 5],
[2, 4],
[3, 1],
[4, 0],
[0, 3]])
基于partition的方法
對于 np.argpartition 函數，復雜度可能下降到 O ( n log K ) O(n \log K)O(nlogK)，很多情況下，K N，此時naive方法有優化的空間。
以下方法首先選出 topK 項，然后僅對前topK項進行排序（matrix僅限2d-array）。
def partition_arg_topK(matrix, K, axis=0):
"""
perform topK based on np.argpartition
:param matrix: to be sorted
:param K: select and sort the top K items
:param axis: 0 or 1. dimension to be sorted.
:return:
"""
a_part = np.argpartition(matrix, K, axis=axis)
if axis == 0:
row_index = np.arange(matrix.shape[1 - axis])
a_sec_argsort_K = np.argsort(matrix[a_part[0:K, :], row_index], axis=axis)
return a_part[0:K, :][a_sec_argsort_K, row_index]
else:
column_index = np.arange(matrix.shape[1 - axis])[:, None]
a_sec_argsort_K = np.argsort(matrix[column_index, a_part[:, 0:K]], axis=axis)
return a_part[:, 0:K][column_index, a_sec_argsort_K]
# Example
>>> dists = np.random.permutation(np.arange(30)).reshape(6, 5)
array([[17, 28, 1, 24, 23, 8],
[ 9, 21, 3, 22, 4, 5],
[19, 12, 26, 11, 13, 27],
[10, 15, 18, 14, 7, 16],
[ 0, 25, 29, 2, 6, 20]])
>>> partition_arg_topK(dists, 2, axis=0)
array([[4, 2, 0, 4, 1, 1],
[1, 3, 1, 2, 4, 0]])
>>> partition_arg_topK(dists, 2, axis=1)
array([[2, 5],
[2, 4],
[3, 1],
[4, 0],
[0, 3]])
大數據量測試
對shape(5000, 100000)的矩陣進行topK排序，測試時間為：
| K |
partition(s) |
naive(s) |
| 10 |
8.884 |
22.604 |
| 100 |
9.012 |
22.458 |
| 1000 |
8.904 |
22.506 |
| 5000 |
11.305 |
22.844 |
補充：python堆排序實現TOPK問題
# 構建小頂堆跳轉def sift(li, low, higt):
tmp = li[low]
i = low
j = 2 * i + 1
while j = higt: # 情況2：i已經是最后一層
if j + 1 = higt and li[j + 1] li[j]: # 右孩子存在并且小于左孩子
j += 1
if tmp > li[j]:
li[i] = li[j]
i = j
j = 2 * i + 1
else:
break # 情況1：j位置比tmp小
li[i] = tmp
def top_k(li, k):
heap = li[0:k]
# 建堆
for i in range(k // 2 - 1, -1, -1):
sift(heap, i, k - 1)
for i in range(k, len(li)):
if li[i] > heap[0]:
heap[0] = li[i]
sift(heap, 0, k - 1)
# 挨個輸出
for i in range(k - 1, -1, -1):
heap[0], heap[i] = heap[i], heap[0]
sift(heap, 0, i - 1)
return heap
li = [0, 8, 6, 2, 4, 9, 1, 4, 6]
print(top_k(li, 3))
以上為個人經驗，希望能給大家一個參考，也希望大家多多支持腳本之家。
您可能感興趣的文章:- Python堆排序原理與實現方法詳解
- Python實現堆排序的方法詳解
- python 的topk算法實例
- python 實現堆排序算法代碼
