We don’t need to add all the elements in the heap but rather check if the diff is smaller and then add the item and remove the other.
from heapq import *
def find_closest_elements(arr, K, X):
index = binary_search(arr, X)
low = max(0, index - K)
high = min(index + K, len(arr) - 1)
max_heap = []
for i in range(low, high + 1):
diff = abs(X - arr[i])
if len(max_heap) < K:
heappush(max_heap, (-diff, arr[i]))
else:
if diff < -max_heap[0][0]:
heappushpop(max_heap, (-diff, arr[i]))
return [heappop(max_heap)[1] for _ in range(K)]
def binary_search(arr, num):
start, end = 0, len(arr) - 1
while start <= end:
mid = (start + end) // 2
if arr[mid] == num:
return mid
elif arr[mid] < num:
start = mid + 1
else:
end = mid - 1
return start
def main():
print("'K' closest numbers to 'X' are: " +
str(find_closest_elements([5, 6, 7, 8, 9], 3, 7)))
print("'K' closest numbers to 'X' are: " +
str(find_closest_elements([2, 4, 5, 6, 9], 3, 6)))
print("'K' closest numbers to 'X' are: " +
str(find_closest_elements([2, 4, 5, 6, 9], 3, 10)))
main()
----
Type your question above this line.
Course: https://www.educative.io/collection/5668639101419520/5671464854355968
Lesson: https://www.educative.io/collection/page/5668639101419520/5671464854355968/5474975244877824