免費(fèi)AI編程助手
topK算法經(jīng)常在各種功能中使用到，在python中，numpy等計(jì)算庫使用了豐富的底層優(yōu)化，對(duì)于矩陣計(jì)算的效率遠(yuǎn)高于python的for-loop實(shí)現(xiàn)。但是numpy中并沒有直接提供topK算法的直接實(shí)現(xiàn)。因此，我們希望盡量用一些numpy函數(shù)的組合實(shí)現(xiàn)topK。接下來這篇文章我們就來介紹一下python怎么實(shí)現(xiàn)topK算法。
pytorch 庫提供了topk函數(shù)，可以將高維數(shù)組沿某一維度（該維度共N項(xiàng)），選出最大（最?。┑腒項(xiàng)并排序。返回排序結(jié)果和index信息。奇怪的是，更輕量級(jí)的numpy庫并沒有直接提供 topK 函數(shù)。numpy只提供了argpartition 和 partition，可以將最大（最小）的K項(xiàng)排到前K位。以argpartition為例，最小的3項(xiàng)排到了前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實(shí)現(xiàn)的是 partial sorting，如上例，前3項(xiàng)和其余項(xiàng)被分開，但是兩部分各自都是不排序的！而我們可能更想要topK的幾項(xiàng)排好序（其余項(xiàng)則不作要求）。因此，下面提供一種基于argpartition的topK方法。
一個(gè)naive方法
最簡單的方法自然是全排序，然后取前K項(xiàng)。缺點(diǎn)在于，要把topK之外的數(shù)據(jù)也進(jìn)行排序，當(dāng)K << N時(shí)較為浪費(fèi)時(shí)間，復(fù)雜度為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的方法
對(duì)于 np.argpartition 函數(shù)，復(fù)雜度可能下降到 O ( n log ? K ) O(n log K)O(nlogK)，很多情況下，K << N，此時(shí)naive方法有優(yōu)化的空間。
以下方法首先選出 topK 項(xiàng)，然后僅對(duì)前topK項(xiàng)進(jìn)行排序（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]])
大數(shù)據(jù)量測(cè)試
對(duì)shape(5000, 100000)的矩陣進(jìn)行topK排序，測(cè)試時(shí)間為：
| 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 |
補(bǔ)充：python堆排序?qū)崿F(xiàn)TOPK問題
# 構(gòu)建小頂堆跳轉(zhuǎn)def sift(li, low, higt):
tmp = li[low]
i = low
j = 2 * i + 1
while j <= higt: # 情況2：i已經(jīng)是最后一層
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)
# 挨個(gè)輸出
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怎么實(shí)現(xiàn)topk算法的全部內(nèi)容了,希望能給大家一個(gè)參考，也希望大家多多支持W3Cschool。
