Top K Frequent Elements Python Solution

lcpy gen -n 347   # by problem number
lcpy gen -s top_k_frequent_elements   # by problem name

Input: nums = [1,1,1,2,2,3], k = 2
Output: [1,2]

Input: nums = [1], k = 1
Output: [1]

Input: nums = [1,2,1,2,1,2,3,1,3,2], k = 2
Output: [1,2]

import heapq
from collections import Counter


class Solution:
    def top_k_frequent(self, nums: list[int], k: int) -> list[int]:
        """
        Optimized version using heap for O(n log k) time complexity.

        Time: O(n log k) - heap operations
        Space: O(n) - for counter and heap
        """
        counter = Counter(nums)

        # Use min heap of size k - keep the k most frequent elements
        heap: list[tuple[int, int]] = []
        for num, count in counter.items():
            if len(heap) < k:
                heapq.heappush(heap, (count, num))
            elif count > heap[0][0]:
                heapq.heapreplace(heap, (count, num))

        # Extract numbers from heap (order doesn't matter for this problem)
        return [num for _, num in heap]