The original question can be found here.
The question is asking to return the k closet elements to target x. The array is sorted.
Heap
First solution is using a heap to keep track the current closest elements to x. Of course the comparator needs to be tweaked accordingly.
class Solution {
public List<Integer> findClosestElements(int[] arr, int k, int x) {
PriorityQueue<Integer> pq = new PriorityQueue<Integer>((a, b) -> {
if (Math.abs(a - x) == Math.abs(b - x)) {
return a >= b ? -1 : 1;
}
return Math.abs(a - x) < Math.abs(b - x) ? 1 : -1;
});
for (int num : arr) {
pq.add(num);
if (pq.size() > k) {
pq.poll();
}
}
List<Integer> res = new ArrayList();
while (pq.size() > 0) {
res.add(pq.poll());
}
Collections.sort(res);
return res;
}
}
We’re using a “max” heap, and the element at root is the one farthest from x. Time complexity: O(nlogk + klogk). Space complexity: O(k). However, this solution does not use leverage the characteristic that array is already sorted.
Binary Search
First we locate the target element in the array, and then we expand to both left and right to find the K closest elements. There’s one case needs to take into consideration: what if we cannot find the element? Actually we don’t have to. We find the first element that is >= target
and expand from there (of course you can choose the last element that is <= target which basically the same idea).
class Solution:
def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
right = self.find_first_greater_or_equal_to_target_element(arr, x)
left = right - 1
res = []
while len(res) < k:
if self.is_left_closer(arr, left, right, x):
res.append(arr[left])
left -= 1
else:
res.append(arr[right])
right += 1
res.sort()
return res
def find_first_greater_or_equal_to_target_element(self, arr: List[int], target: int) -> int:
left, right = 0, len(arr) - 1
while left + 1 < right:
mid = (left + right) // 2
if arr[mid] < target:
left = mid
else:
right = mid
if arr[right] >= target:
return right
if arr[left] >= target:
return left
return len(arr)
def is_left_closer(self, arr: List[int], left: int, right: int, target: int) -> bool:
if left < 0:
return False
if right >= len(arr):
return True
return abs(arr[left] - target) <= abs(arr[right] - target)
Time Complexity: O(logn + k). Space Complexity: O(1).