fork download
  1. import heapq
  2.  
  3. def getGreatestElements(arr, k):
  4. heap = []
  5. ans = []
  6.  
  7. for i, val in enumerate(arr):
  8. heapq.heappush(heap, val)
  9. if len(heap) > k:
  10. heapq.heappop(heap)
  11. if i + 1 >= k:
  12. ans.append(heap[0])
  13.  
  14. return ans
  15.  
  16. # Example:
  17. # n = 4
  18. # k = 2
  19. # arr = [4, 2, 1, 3]
  20. # n = 5
  21. # k = 4
  22. # arr = [3,2,4,5,1]
  23. n= 5
  24. k=2
  25. arr = [3,1,4,5,2]
  26. print(getGreatestElements(arr, k)) # Output: [2, 2, 3]
  27.  
Success #stdin #stdout 0.1s 14120KB
stdin
Standard input is empty
stdout
[1, 3, 4, 4]