Top K Frequent Elements

import heapq
from collections import Counter


# Approach 1: Counter + sort — O(n log n)
def top_k_counter(nums: list[int], k: int) -> list[int]:
    return [x for x, _ in Counter(nums).most_common(k)]


# Approach 2: Min-heap of size k — O(n log k)
def top_k_heap(nums: list[int], k: int) -> list[int]:
    freq = {}
    for n in nums:
        freq[n] = freq.get(n, 0) + 1

    heap = []
    for num, count in freq.items():
        heapq.heappush(heap, (count, num))
        if len(heap) > k:
            heapq.heappop(heap)

    return [num for _, num in heap]


# Approach 3: Bucket sort — O(n) optimal
def top_k_bucket(nums: list[int], k: int) -> list[int]:
    freq = {}
    for n in nums:
        freq[n] = freq.get(n, 0) + 1

    buckets = [[] for _ in range(len(nums) + 1)]
    for num, count in freq.items():
        buckets[count].append(num)

    result = []
    for i in range(len(buckets) - 1, -1, -1):
        for num in buckets[i]:
            result.append(num)
            if len(result) == k:
                return result


if __name__ == "__main__":
    nums = [1, 1, 1, 2, 2, 3]
    print(top_k_counter(nums, 2))  # [1, 2]
    print(top_k_heap(nums, 2))     # [1, 2]
    print(top_k_bucket(nums, 2))   # [1, 2]