-
Notifications
You must be signed in to change notification settings - Fork 0
/
347.前-k-个高频元素.py
58 lines (48 loc) · 1.88 KB
/
347.前-k-个高频元素.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
#
# @lc app=leetcode.cn id=347 lang=python3
#
# [347] 前 K 个高频元素
#
# @lc code=start
# Tips: array, hash-table, heap
class Solution:
def topKFrequent(self, nums: List[int], k: int) -> List[int]:
# 调整堆的函数
def adjust_heap(array, start_pos, end_pos):
temp = array[start_pos]
pos = start_pos
child_pos = pos * 2 + 1
while child_pos <= end_pos:
right_pos = child_pos + 1
if right_pos <= end_pos and array[right_pos][1] > array[child_pos][1]:
child_pos = right_pos
if temp[1] < array[child_pos][1]:
array[pos] = array[child_pos]
pos = child_pos
child_pos = 2 * pos + 1
else:
break
# 注意这里是 pos,不是 child_pos
array[pos] = temp
return
# 计数变成哈希表
count_map = {}
for num in nums:
count_map[num] = count_map.get(num, 0) + 1
# 把哈希表转成数组方便排序
count_list = [(k, v) for k, v in count_map.items()]
# 建立大根堆
length = len(count_list)
for i in range(length // 2, -1, -1):
adjust_heap(count_list, i, length - 1)
# 调整堆,每次选出一个最大的,循环 k 次
ans = []
for i in range(length - 1, length - 1 - k, -1):
# 现在最大的是 count_list[0],和待排序的最后一个交换
count_list[0], count_list[i] = count_list[i], count_list[0]
# 此时 i 已经在排好序后应该在的位置了,需要排序的是 0 到 i-1
adjust_heap(count_list, 0, i - 1)
# 把第 i 个已经排好位置的加入到结果列表里
ans.append(count_list[i][0])
return ans
# @lc code=end