40 lines
688 B
Python
40 lines
688 B
Python
#!/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))
|