leetcode-347. 前 K 个高频元素

题目描述

给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。

示例 1:

输入: nums = [1,1,1,2,2,3], k = 2
输出: [1,2]
示例 2:

输入: nums = [1], k = 1
输出: [1]

思路

1）先用字典保存每个数出现的次数

2）用小根堆，如果当前小根堆保留长度小于k个，刚进来的元素直接放进去就可以

3）否则，比较进来元素的次数如果大于堆顶的频率的话，就先将刚进来的元素放进去，之后pop栈顶元素（heapq.heappushpop：先放进，再pop）

import heapq


class Solution(object):
    def topKFrequent(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[int]
        """
        freq_dict = {}
        for num in nums:
            if num not in freq_dict:
                freq_dict[num] = 1
            else:
                freq_dict[num] += 1
        min_heap = []
        for i, v in freq_dict.items():
            if len(min_heap) < k:
                heapq.heappush(min_heap, (v, i))
            elif v > min_heap[0][0]:
                heapq.heappushpop(min_heap, (v, i))
        res = [i for v, i in min_heap]
        return res

if __name__=='__main__':
    s=Solution()
    nums = [1, 1, 1, 2, 2, 3]
    k = 2
    print(s.topKFrequent(nums, k))

二刷：

import heapq
class Solution(object):
    def topKFrequent(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: List[int]
        """
        c_dict = {}
        for num in nums:
            if num not in c_dict:
                c_dict[num] = 1
            else:
                c_dict[num]+=1
        min_heap = []
        for key, value in c_dict.items():
            if len(min_heap)<k:
                heapq.heappush(min_heap, (value, key))
            elif value > min_heap[0][0]:
                heapq.heappushpop(min_heap, (value,key))
        res = [i[1] for i in min_heap]
        return res 

 记忆点：

小根堆

记住写法，无论push,pop，都要加heap