主頁 > 知識庫 > 使用numpy實現(xiàn)topk函數(shù)操作(并排序)

使用numpy實現(xiàn)topk函數(shù)操作(并排序)

熱門標簽:江西省地圖標注 沈陽外呼系統(tǒng)呼叫系統(tǒng) 外呼系統(tǒng)哪些好辦 沈陽人工外呼系統(tǒng)價格 如何申請400電話費用 武漢外呼系統(tǒng)平臺 富錦商家地圖標注 沈陽防封電銷卡品牌 池州外呼調(diào)研線路

np.argpartition 難以解決topK

topK是常用的一個功能,在python中,numpy等計算庫使用了豐富的底層優(yōu)化,對于矩陣計算的效率遠高于python的for-loop實現(xiàn)。因此,我們希望盡量用一些numpy函數(shù)的組合實現(xiàn)topK。

pytorch 庫提供了topk函數(shù),可以將高維數(shù)組沿某一維度(該維度共N項),選出最大(最小)的K項并排序。返回排序結(jié)果和index信息。奇怪的是,更輕量級的numpy庫并沒有直接提供 topK 函數(shù)。numpy只提供了argpartition 和 partition,可以將最大(最?。┑腒項排到前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實現(xiàn)的是 partial sorting,如上例,前3項和其余項被分開,但是兩部分各自都是不排序的!而我們可能更想要topK的幾項排好序(其余項則不作要求)。因此,下面提供一種基于argpartition的topK方法。

一個naive方法

最簡單的方法自然是全排序,然后取前K項。缺點在于,要把topK之外的數(shù)據(jù)也進行排序,當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 函數(shù),復雜度可能下降到 O ( n log ⁡ K ) O(n \log K)O(nlogK),很多情況下,K N,此時naive方法有優(yōu)化的空間。

以下方法首先選出 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]])

大數(shù)據(jù)量測試

對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堆排序?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)
    # 挨個輸出
    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))

以上為個人經(jīng)驗,希望能給大家一個參考,也希望大家多多支持腳本之家。

您可能感興趣的文章:
  • Python堆排序原理與實現(xiàn)方法詳解
  • Python實現(xiàn)堆排序的方法詳解
  • python 的topk算法實例
  • python 實現(xiàn)堆排序算法代碼

標簽:阿里 潛江 通遼 黑龍江 銅川 常德 株洲 呂梁

巨人網(wǎng)絡通訊聲明:本文標題《使用numpy實現(xiàn)topk函數(shù)操作(并排序)》,本文關(guān)鍵詞  使用,numpy,實現(xiàn),topk,函數(shù),;如發(fā)現(xiàn)本文內(nèi)容存在版權(quán)問題,煩請?zhí)峁┫嚓P(guān)信息告之我們,我們將及時溝通與處理。本站內(nèi)容系統(tǒng)采集于網(wǎng)絡,涉及言論、版權(quán)與本站無關(guān)。
  • 相關(guān)文章
  • 下面列出與本文章《使用numpy實現(xiàn)topk函數(shù)操作(并排序)》相關(guān)的同類信息!
  • 本頁收集關(guān)于使用numpy實現(xiàn)topk函數(shù)操作(并排序)的相關(guān)信息資訊供網(wǎng)民參考!
  • 推薦文章