DEV Community

Riyaz
Riyaz

Posted on

'K' Closest Numbers(code and explanation)

Problem Statement
Given a number array and two integers ‘K’ and ‘X’, find ‘K’ closest numbers to ‘X’ in the array.
Return the numbers in the sorted order. ‘X’ is not necessarily present in the array.
Example 1:
Input: [5, 6, 7, 8, 9], K = 3, X = 7
Output: [6, 7, 8]
Example 2:
Input: [2, 4, 5, 6, 9], K = 3, X = 6
Output: [4, 5, 6]
SOLUTION
For finding the closest numbers to X we can simple calculate the absolute difference between X and
all the values from the array.Now we store these values in a pair using tuple (abs diff, value) with the abs difference
along with the value itself
So, for example, [5, 6, 7, 8, 9], K = 3, X = 7
the tuple values would be (2,5), (1,6), (0,7), (1,8), (2,9)
So the value the first value in the tuple, the closest the element is from X
Now we want the top k(3) values from this tuple that are closest to X(7)
Now we can utilize the heap data structure to do so
We can use min heap which will store the tuples in smallest to greatest values
So the heap will look like

            (0,7)
           /     \
         (1,6)   (1,8)
        /    \
      (2,5)  (2,9)     
Enter fullscreen mode Exit fullscreen mode

Now we can easily remove the top K elements from this heap which will give the closest numbers
Look at the below code for more clarity
CODE

from heapq import *
def findClosestElements(A,K,X):
    answer = []

    minHeap1 = []
    for num in A:
        heappush(minHeap1, (abs(num-X),num)) #Creates a min heap with tuple values as discussed

    minHeap2 = [] #Remove the top K elements from minHeap1 and store in minHeap2 as we have to maintain order
    n = K
    while n > 0:
        a = heappop(minHeap1)
        heappush(minHeap2, a[1]) #Append top K elements values from the minHeap1 and store the second value 
        n -= 1                   # from the tuple i.e (0,7) --> 7 as '7' is the actual element

    while minHeap2:
        b = heappop(minHeap2) #Remove all the elements from the minHeap2 one by one and store them in answer 
        answer.append(b)

    return answer 
Enter fullscreen mode Exit fullscreen mode

Time Complexity
The first step of creating minHeap1 will take O(N)
Creating minHeap2 will take O(KlogK)
So overall time complexity is O(KLogK)
Space Complexity
O(N) for storing N elements in heap

Top comments (0)