Leetcode 658 Solution

This article provides solution to leetcode question 658 (find-k-closest-elements)

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

Solution

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        l = 0
        r = len(arr) - 1
        while l < r:
            m = (l + r) // 2
            if arr[m] >= x:
                r = m
            else:
                l = m + 1

        r = l
        l = l - 1

        ans1 = []
        ans2 = []
        while (l > 0 or r < len(arr)) and len(ans1) + len(ans2) < k:
            left_dist = x - arr[l] if l >= 0 else 100000
            right_dist = arr[r] - x if r < len(arr) else 100000

            if left_dist <= right_dist:
                ans1.append(arr[l])
                l -= 1
            else:
                ans2.append(arr[r])
                r += 1

        return list(reversed(ans1)) + ans2