【Python】快速选择法：Leetcode：347. 前 K 个高频元素模板做题就是爽

原创于 2024-07-03 14:59:00 发布 · 287 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#python #leetcode #算法 #深度优先 #排序算法

描述

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

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

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

提示：
1 <= nums.length <= 105
k 的取值范围是 [1, 数组中不相同的元素的个数]
题目数据保证答案唯一，换句话说，数组中前 k 个高频元素的集合是唯一的

进阶：你所设计算法的时间复杂度必须优于 O(n log n) ，其中 n 是数组大小。

思路

直接套用快速选择模板，详见我上一篇文章，时间复杂度O(N)

代码

class Solution:
    def topKFrequent(self, nums: List[int], k: int) -> List[int]:
        self.fre_map = {}
        for num in nums:
            if(num not in self.fre_map): self.fre_map[num] = 0
            self.fre_map[num] += 1
        
        self.keys = list(self.fre_map.keys())
        self.quickSchedule(0,len(self.keys)-1,k-1)
        return self.keys[:k]
    
    def quickSchedule(self,l,r,k):
        if(l==r): return
        pivot = self.keys[l]
        i = l-1
        j = r+1
        while(i < j):
            i += 1
            while(self.fre_map[self.keys[i]] > self.fre_map[pivot]): i += 1
            j -= 1
            while(self.fre_map[self.keys[j]] < self.fre_map[pivot]): j -= 1
            if(i < j): self.keys[i],self.keys[j] = self.keys[j],self.keys[i]
        
        if(j > k):
            self.quickSchedule(l,j,k)
        else:
            self.quickSchedule(j+1,r,k)