Find K Closest Elements – Python 3 (Week 17 – 15)

class Solution:
    """
    @param A: an integer array
    @param target: An integer
    @param k: An integer
    @return: an integer array
    """
    def kClosestNumbers(self, nums, target, k):
        # write your code here
        if not nums or not target or not k:
            return []
            
        if k > len(nums):
            return nums
            
        start = 0
        end = len(nums) - 1
        
        while start + 1 < end:
            mid = start + (end - start) // 2
            if nums[mid] < target:
                start = mid
            else:
                end = mid
                
        result = []
        
        left, right = start, end
        
        for _ in range(k):
            if left < 0:
                result.append(nums[right])
                right += 1
            elif right >= len(nums):
                result.append(nums[left])
                left -= 1
            elif target - nums[left] <= nums[right] - target:
                result.append(nums[left])
                left -= 1
            else:
                result.append(nums[right])
                right += 1
        
        return result
                
            

Time O(logn +k). Space O(1).

Leave a comment

Design a site like this with WordPress.com
Get started