forked from wangzheng0822/algo
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtop_k.py
More file actions
39 lines (29 loc) · 688 Bytes
/
top_k.py
File metadata and controls
39 lines (29 loc) · 688 Bytes
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
#!/usr/bin/python
# -*- coding: UTF-8 -*-
import random
from heap import MinHeap
def top_k(nums, k):
"""
返回数组的前k大元素
:param nums:
:param k:
:return:
"""
if len(nums) <= k:
return nums
min_h = MinHeap(nums[:k], k)
for i in range(k, len(nums)):
tmp = min_h.get_top()
if nums[i] > tmp:
min_h.remove_top()
min_h.insert(nums[i])
return min_h.get_data()
if __name__ == '__main__':
nums = []
k = 3
for i in range(20):
nums.append(random.randint(1, 100))
print('--- nums ---')
print(nums)
print('--- top {} ---'.format(k))
print(top_k(nums, k))