k_closest_element.py

import heapq


class Solution:
    def findClosestElements(self, arr, k, x):
        if not arr:
            return []

        def calculate_difference(num):
            return abs(num - x)

        heap = []  # max-heap (-difference, num)

        for num in arr:
            difference = -calculate_difference(num)
            if len(heap) < k:
                heapq.heappush(heap, (difference, num))
            elif difference > heap[0][0]:
                heapq.heappushpop(heap, (difference, num))

        return sorted([num for _, num in heap])

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


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

    arr = [1, 2, 3, 4, 5]
    k = 4
    x = 3
    result = solution.findClosestElements(arr, k, x)
    assert result == [1, 2, 3, 4]
    print(result)
    print("Test Cased 1 Passed!")