Problem statement

https://leetcode.com/problems/find-k-closest-elements/

Solution

Let us first solve given problem when k = 1, that is we need to find closest element in our tree. For this we can use binary search and then check two possible options for element just before x and for element just after x. Once we found the closest number, we can have two possible options: either the next closest number is smaller or bigger, which one to choose? We check the distance to both of them and choose the closest one. On the next step we also have two choices and so on. So, what we have in the end is two pointers approach and on each step we decide which element we take.

Complexity

Time complexity is O(log n + k), space complexity is O(k).

Code

from bisect import bisect, bisect_left

class Solution:
    def findClosestElements(self, arr, k, x):
        n = len(arr)
        ind = bisect_left(arr, x)
        if ind == n or ind > 0 and arr[ind] + arr[ind-1] >= 2*x:
            ind -= 1

        beg, end = ind, ind
        for i in range(k-1):
            if beg == 0: end += 1
            elif end == n-1: beg -= 1
            else:
                if arr[end+1] + arr[beg-1] >= 2*x:
                    beg -= 1
                else:
                    end += 1

        return arr[beg:end+1]