力扣 347. 前K个高频元素

解法: 堆

def topKFrequent(nums, k):
    def heap_up(i):
        cur, parent = i, (i - 1) >> 1
        while cur > 0 and heap[cur][1] >= heap[parent][1]:
            heap[cur], heap[parent] = heap[parent], heap[cur]
            heap_down(cur)
            cur = parent
            parent = (cur - 1) >> 1
        
    def heap_down(i):
        cur, left, right = i, 2 * i + 1, 2 * i + 2
        while left < heap_num :
            x, nex = heap[cur][1], cur
            if left < heap_num and x < heap[left][1]:
                nex = left
                x = heap[left][1]
            if right < heap_num and x < heap[right][1]:
                nex = right
            if cur != nex:
                heap[cur], heap[nex] = heap[nex], heap[cur]
                cur = nex
                left, right = 2 * cur + 1, 2 * cur + 2
            else:
                break
            

    # 堆排序
    heap = []
    heap_num = 0
    for num in nums:
        i = 0
        while i < heap_num and heap[i][0] != num: i += 1
        if i >= heap_num: 
            heap.append([num, 1])
            heap_num += 1
        else:
            heap[i][1] += 1
            # 上浮堆元素
            heap_up(i)
    # print(heap)
    # 输出结果
    ret = []
    for i in range(k):
        ret.append(heap[0][0])
        heap[0], heap[-1] = heap[-1], heap[0]
        heap.pop()
        heap_num -= 1
        heap_down(0)
        # print(ret[-1], heap)
    # print(ret)
    return ret