k_most_frequent.py

#!/usr/bin/env python
import sys

# TestCase = (nums, k, possible_results)
type TestCase = tuple[list[int], int, list[set[int]]]


def k_most_frequent(nums: list[int], k: int) -> set[int]:
    """
    Returns a set of k most frequently recurring unique numbers.
    If k > uniquely frequent elements or len(nums), we just return
    all the unique numbers - no need to change that; confirmed.
    """

    if k <= 0:
        return set()

    result: set[int] = set()
    counts: dict[int, int] = {}
    buckets: list[list[int]] = [[] for _ in range(len(nums) + 1)]

    # Slightly faster than `counts[num] = counts.get(num, 0) + 1` due
    # to byte-code optimization paths in modern Python interpreters.
    for num in nums:
        try:
            counts[num] += 1
        except KeyError:
            counts[num] = 1

    for num, count in counts.items():
        buckets[count].append(num)

    for bucket_idx in range(len(buckets) - 1, 0, -1):
        for num in buckets[bucket_idx]:
            result.add(num)

            if len(result) == k:
                return result

    return result


def run_test(test_case: TestCase) -> bool:
    """
    Runs a test case and returns whether the test was successful or not.
    """

    test_nums, test_k, expected = test_case
    actual = k_most_frequent(test_nums, test_k)

    if actual not in expected:
        print(f'❌ failure for {test_nums}; {expected=} | {actual=}', file=sys.stderr)
        return False

    print('✅ success')
    return True


def main() -> int:

    tests: list[TestCase] = [
        ([1], 1, [{1}]),
        ([], 0, [set()]),
        ([1,2], 5, [{1,2}]),
        ([1,1,2], 0, [set()]),
        ([1,1,1,2,2,3], 2, [{1, 2}]),
        ([1,2,3,4,4,5,5,6,6,6], 3, [{4, 5, 6}]),
        ([1,1,1,1,2,2,2,2,3,3,3,3], 2, [{1,2}, {2, 3}, {1, 3}]),
        ([1,1,1,1,1,2,2,2,2,3,3,3,3,3,3,3,3,4,4,5,6,6,6,6,6,7], 3, [{3,6,1}])
    ]

    if not all(run_test(test) for test in tests):
        return -1

    return 0


if __name__ == '__main__':
    sys.exit(main())