class Solution:
def partition(self, nums, left, right):
"""划分函数:以nums[right]为基准,返回基准值的正确位置索引"""
pivot = nums[right]
i = left - 1 # 指向小于基准的子数组末尾
for j in range(left, right):
if nums[j] <= pivot:
i += 1
nums[i], nums[j] = nums[j], nums[i] # 将小元素交换到左侧
nums[i+1], nums[right] = nums[right], nums[i+1] # 基准归位
return i + 1
def topk_split(self, nums, k, left, right):
"""快速选择算法核心:找到第k小元素的位置后停止递归"""
if left < right:
# 注意必须添加 self. 调用类方法
index = self.partition(nums, left, right)
if index == k:
return # 找到目标位置,终止递归
elif index < k:
self.topk_split(nums, k, index+1, right) # 处理右半部分
else:
self.topk_split(nums, k, left, index-1) # 处理左半部分
#获得前k小的数
def topk_smalls(nums, k):
topk_split(nums, k, 0, len(nums)-1)
return nums[:k]
#获得前k大的数
def topk_larges(nums, k):
#parttion是按从小到大划分的,如果让index左边为前n-k个小的数,则index右边为前k个大的数
topk_split(nums, len(nums)-k, 0, len(nums)-1) #把k换成len(nums)-k
return nums[len(nums)-k:]
一、215. 数组中的第K个最大元素
思路:经典的快排和堆的题目
代码1:快排(这道题单纯使用快排无法ak)
python复制代码
class Solution:
def findKthLargest(self, nums, k):
def quick_sort(nums, k):
pivot = random.choice(nums)
big, equal, small = [], [], []
for num in nums:
if num > pivot:
big.append(num)
elif num < pivot:
small.append(num)
else:
equal.append(num)
if k <= len(big): # 第 k 大元素在 big 中,递归划分
return quick_sort(big, k)
if k > len(nums)-len(small): # 第 k 大元素在 small 中,递归划分
return quick_sort(small, k-len(nums)+len(small))
return pivot # 第 k 大元素在 equal 中,直接返回 pivot
return quick_sort(nums, k)
代码2:堆
python复制代码
class Solution:
def findKthLargest(self, nums: List[int], k: int) -> int:
pq = [] # 将数组加入小顶堆,堆中维护当前值最大的k个数
for num in nums:
heapq.heappush(pq, num) # 当前元素入堆
if len(pq) > k:
heapq.heappop(pq) # 堆中元素超过k个,弹出最小的那个
return pq[0] # 最后堆顶的即为第k大的数
二、347. 前 K 个高频元素
思路:
用字典保存每个数字出现次数,最后排序key
代码:
python复制代码
class Solution:
def topKFrequent(self, nums: List[int], k: int) -> List[int]:
# 使用 defaultdict 统计每个元素出现的次数
coll_dea = collections.defaultdict(int)
for i in nums:
coll_dea[i] += 1
# 按频率排序(从高到低),然后取前 k 个元素
sorted_items = sorted(coll_dea.items(), key=lambda x: x[1], reverse=True)
ans = [item[0] for item in sorted_items[:k]] # 提取前 k 个元素的值
return ans