题目链接

题解

1. 前 K 个高频元素

思路

代码

class Solution:
    ### 0128 堆排序（60 ms，17.5 MB）
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        def heapify(p_idx, heap_size):
            max_idx = p_idx
            l_idx, r_idx = p_idx*2+1, p_idx*2+2

            # 双元素版本的建堆，因此需要对比对应位置的元素（这里的第二个下标是1）
            if l_idx < heap_size and new_arr[l_idx][1] > new_arr[max_idx][1]:
                max_idx = l_idx
            
            if r_idx < heap_size and new_arr[r_idx][1] > new_arr[max_idx][1]:
                max_idx = r_idx

            # 交换不变（整体交换）
            if max_idx != p_idx:
                new_arr[max_idx], new_arr[p_idx] = new_arr[p_idx], new_arr[max_idx]
                heapify(max_idx, heap_size)

        # 先遍历数组，统计每个元素出现的数量，然后转化成新的双元素数组new_arr
        stat = collections.Counter(nums)
        new_arr = list(stat.items())
        n = len(new_arr) # 注意：这里的n是新数组的长度！

        res = []

        for i in range(n//2-1, -1, -1):
            heapify(i, n)
        
        # 这里只需维护k个最大值，因此到n-k-1（否则到0）
        for heap_size in range(n-1, n-k-1, -1):
            # 记录每次得到的最大根（这里的第二个下标是0）
            res.append(new_arr[0][0])

            # 交换不变（整体交换）
            new_arr[0], new_arr[heap_size] = new_arr[heap_size], new_arr[0]
            heapify(0, heap_size)
        
        return res