kth_largest.py

import heapq


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

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

        return heap[0]

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

    def kthLargest(self, nums, k):
        nums = [-num for num in nums]
        heapq.heapify(nums)

        for _ in range(k - 1):
            heapq.heappop(nums)

        return -heapq.heappop(nums)

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


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

    nums = [5, 3, 2, 1, 4]
    k = 2
    result = solution.kthLargest(nums, 2)
    assert result == 4
    print(result)
    print("Test Case 1 Passed!")