See all coding puzzles and patterns here.

Problem Definition

Given an array of positive integers and a target integer, return the $k$ elements closest to the target. Assume that all integers in the array are unique.

Examples

Some test cases to verify our solution.

assert closest_k_elements([1, 2, 3, 4, 5], 10, 500) == [1, 2, 3, 4, 5]
assert closest_k_elements([1, 2, 3, 4, 5], 3, 3) == [2, 3, 4]
assert closest_k_elements([1, 2, 3, 4, 5], 4, 3) == [3, 4, 5]
assert closest_k_elements([1, 2, 3, 4, 8, 12, 13, 14, 15], 7, 3) == [3, 4, 8]

Solution

We use a binary search to find integer closest to the target and then do a merge of the items to the left and right of it until we have found the $k$ closest elements. 

from collections import deque

def find_closest_idx(elements, target):
    left = 0
    right = len(elements)
    idx = right // 2
    while 1:
        if target > elements[idx]:
            left = idx
            next_idx = idx + (right - idx) // 2
        elif target < elements[idx]:
            right = idx
            next_idx = idx - (idx - left) // 2
        elif target == elements[idx]:
            break
        if next_idx <= left or next_idx >= right:
            break
        idx = next_idx
    return idx

def find_k_nearest(elements, target, closest_idx, k):
    result = deque()
    left = closest_idx
    right = closest_idx + 1
    elements_len = len(elements)
    while len(result) < k:
        if left < 0:
            result.extend(elements[right:right+k-len(result)])
        elif right >= elements_len:
            result.extendleft(elements[left-k+len(result)+1:left+1])
        elif abs(target - elements[left]) < abs(elements[right] - target):
            result.appendleft(elements[left])
            left -= 1
        else:
            result.append(elements[right])
            right += 1
    return list(result)

def closest_k_elements(elements, target, k):
    if len(elements) <= k:
        return elements
    if target < elements[0]:
        return elements[:k]
    if target > elements[-1]:
        return elements[len(elements) - k:]
    return find_k_nearest(elements, target, find_closest_idx(elements, target), k)

Complexity

Let $n = $ len(elements).

The runtime complexity of find_closest_idx should be $O(n \cdot log(n))$.
And its auxiliary space complixity should be $O(1)$, since we are only storing some contants.

The runtime complexity of find_k_nearest should be $\Theta(k)$ since we will look at $k+1$ items.
And its auxiliary space complexity should also be $\Theta(1)$ since we are only retaining constant-sized variables.
The result array is the return value; so while its space complexity is $\Theta(k)$, it is not auxiliary space.

The runtime complexity of closest_k_elements should be $O(k + log(n))$. And its auxiliary space complexity should be $O(1)$.