Python 解决Topk问题

---

Topk问题，即是在长度为n的列表中找到最大的（或最小的）k个。

排序后切片

时间复杂度 O(nlogn)

def topk(arr, k):
    arr.sort(reverse=True)
    ret = arr[0: k]
    return ret

局部排序

采用类似冒泡排序的思想，进行局部排序，外层循环只进行k次，每次从无序部分找到最大的一个数。

时间复杂度 O(kn)

def topk(arr, k):
    for i in range(k):
        for j in range(len(arr) - 1, i, -1):
            if arr[j] > arr[j - 1]:
                arr[j], arr[j - 1] = arr[j - 1], arr[j]
    return arr[0: k]

堆

维护一个小根堆，堆的大小为k，遍历数组判断是否可以入堆，再将堆逐个出数。

时间复杂度 O(klogn)

def shift(arr, low, high):
    i = low
    j = 2 * i + 1
    temp = arr[i]
    while j <= high:
        if j + 1 <= high and arr[j + 1] < arr[j]:
            j += 1
        if arr[j] < temp:
            arr[i] = arr[j]
            i = j
            j = 2 * i + 1
        else:
            break
    arr[i] = temp


def topk(arr, k):
    # 建堆
    heap = arr[0: k]
    for i in range((k - 2) // 2, -1, -1):
        shift(heap, i, k - 1)
    # 遍历
    for i in range(k, len(arr) - 1):
        if arr[i] > heap[0]:
            heap[0] = arr[i]
            shift(heap, 0, k - 1)
    # 出数
    for i in range(k - 1, -1, -1):
        heap[0], heap[i] = heap[i], heap[0]
        shift(heap, 0, i - 1)
    return heap