658. Find K Closest Elements

py

import bisect

class Solution:
    def findClosestElements(self, arr, k, x):
        """
        :type arr: List[int]
        :type k: int
        :type x: int
        :rtype: List[int]
        """
        pos = bisect.bisect_left(arr, x)
        # fix pos to get the most close element
        if pos > 0:
            pos = min([pos, pos-1], key=lambda i: abs(x-arr[i]))
        res = [arr[pos]]
        left = pos-1
        right = pos+1
        while len(res) < k:
            # print(res, left, right)
            if left > -1 and right < len(arr):
                # both pos ok, compare and choose one
                small = arr[left]
                big = arr[right]
                if abs(big - x) < abs(small - x):
                    # choose big
                    res.append(big)
                    right += 1
                else:
                    # choose small
                    res.insert(0, small)
                    left -= 1
            elif left > -1:
                # only left ok
                # choose small
                res.insert(0, arr[left])
                left -= 1
            elif right < len(arr):
                # only right ok
                # choose big
                res.append(arr[right])
                right += 1
            else:
                # neither ok
                return res
        return res