Top K Frequent Elements — топ-K самых частых элементов

Условие

Дан массив nums и число k. Верните k самых часто встречающихся элементов в любом порядке.

O(n log n) неприемлемо — нужно лучше.

Решение

Подход 1: heap (O(n log k))

Считаем частоты, кладём в min-heap размера k.

class Solution:
    def topKFrequent(self, nums: list[int], k: int) -> list[int]:
        import heapq
        from collections import Counter
        cnt = Counter(nums)
        heap = []
        for num, freq in cnt.items():
            heapq.heappush(heap, (freq, num))
            if len(heap) > k:
                heapq.heappop(heap)
        return [num for _, num in heap]

Подход 2: bucket sort (O(n))

Частоты лежат в [1..n]. Заводим n + 1 «вёдер»; в bucket[freq] кладём все числа с такой частотой. Идём по bucket от n к 0, набираем k чисел.

def topKFrequent(nums, k):
    from collections import Counter
    cnt = Counter(nums)
    buckets: list[list[int]] = [[] for _ in range(len(nums) + 1)]
    for num, f in cnt.items():
        buckets[f].append(num)
    res = []
    for f in range(len(buckets) - 1, 0, -1):
        for num in buckets[f]:
            res.append(num)
            if len(res) == k:
                return res
    return res

Подход 3: Quickselect (O(n) среднее)

По частотам — выбрать k-ю по убыванию через quickselect. Чуть сложнее, но O(n) в среднем.

Сложность

• Подход 1: O(n log k) времени, O(n + k) памяти.
• Подход 2: O(n) времени, O(n) памяти.

Подводные камни

1. Сортировать по частоте. sorted(cnt.items(), key=...)[:k] — это O(n log n). На собесе требуют лучше.
2. heapq.nlargest. Под капотом тот же O(n log k), на собесе обычно ожидают «руками».
3. Стабильность порядка. Условие позволяет любой порядок — это упрощает.

Эталонный ответ

Bucket sort по частотам — O(n). Heap-вариант — O(n log k), проще для написания.