educative.io

Educative

We can further optimize and reduce the space and time complexity

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