top_k_largest.py

import heapq


class Solution:
    def topKLargest(self, nums):
        heap = []
        k = 3

        """
        If the new element is more frequent, it deserves to be in the top K, so you pop the smallest frequency and push the new one.

        """

        for num in nums:
            if len(heap) < k:
                heapq.heappush(heap, num)
            elif num > heap[0]:
                heapq.heappushpop(heap, num)

        return heap

        # Time Complexity: O(n log k)
        # Space Complexity: O(k)

    def topKLargestLessEfficient(self, nums):
        # top 3 largest
        # not efficient
        nums = [-num for num in nums]
        heapq.heapify(nums)

        k = 3
        result = []

        for _ in range(k - 1):
            max_val = -heapq.heappop(nums)
            result.append(max_val)

        result.append(-heapq.heappop(nums))
        return result

        # Time Complexity: O(n + n log k)
        # Space Complexity: O(n)


if __name__ == "__main__":
    solution = Solution()

    nums = [9, 3, 7, 1, -2, 6, 8]
    result = solution.topKLargest(nums)
    assert set(result) == {7, 8, 9}
    print(result)
    print("Test Case 1 Passed!")