pylist

460. Find K Closest Elements

Given target, a non-negative integer k and an integer array A sorted in ascending order, find the k closest numbers to target in A, sorted in ascending order by the difference between the number and target. Otherwise, sorted in ascending order by number if the difference is same.

Solution

class Solution:
    """
    @param A: an integer array
    @param target: An integer
    @param k: An integer
    @return: an integer array
    """
    def kClosestNumbers(self, A, target, k):
        # init 
        up = len(A)
        low = 0
        # binary search
        while low + 1 < up:
            mid = (low + up) // 2
            if A[mid] > target:
                up = mid
            if A[mid] < target:
                low = mid
            if A[mid] == target:
                up = mid + 1
                low = mid 
                break
        # use up and low as kappa pointers
        ans = []
        while len(ans) < k:
            # if one pointer cannot move, move the other
            if not self.inbond(A, low):
                ans.append(A[up])
                up += 1
                continue
            if not self.inbond(A, up):
                ans.append(A[low])
                low -= 1
                continue
            # move the closer one
            if target - A[low] <= A[up] - target:
                ans.append(A[low])
                low -= 1
            else:
                ans.append(A[up])
                up += 1
        return ans

    def inbond(self, A, i):
        if i < 0 or i >= len(A):
            return False
        return True

special care