Find K Closest Elements

NeetCode: Find K Closest Elements · LeetCode: 658. Find K Closest Elements · Difficulty: Medium

Given a sorted integer array arr, an integer k and an integer x, return the k closest integers to x in the array, sorted in ascending order. An integer a is closer to x than b if |a - x| < |b - x|, or if the distances are equal and a < b.

Example

Text

Input:  arr = [1, 2, 3, 4, 5], k = 4, x = 3      Output: [1, 2, 3, 4]
Input:  arr = [1, 1, 2, 3, 4, 5], k = 4, x = -1  Output: [1, 1, 2, 3]

Brute force: sort by distance

Copy the array into a List<Integer> and sort it with a comparator on the pair (|a - x|, a), which is exactly the "closer" order defined above. Take the first k and sort those back into ascending order. This ignores that the input is already sorted. Distances are computed in long because a - x can overflow int for extreme values. The comparator uses Long.compare and Integer.compare instead of subtraction, which can overflow.

Java

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        List<Integer> v = new ArrayList<>();
        for (int a : arr) {
            v.add(a);
        }
        v.sort((a, b) -> {
            int cmp = Long.compare(Math.abs((long) a - x), Math.abs((long) b - x));
            return cmp != 0 ? cmp : Integer.compare(a, b);
        });
        List<Integer> res = new ArrayList<>(v.subList(0, k));
        Collections.sort(res);
        return res;
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.findClosestElements(new int[] {1, 2, 3, 4, 5}, 4, 3));
        System.out.println(s.findClosestElements(new int[] {1, 1, 2, 3, 4, 5}, 4, -1));
        System.out.println(s.findClosestElements(new int[] {1, 3}, 2, 2));
        System.out.println(s.findClosestElements(new int[] {5}, 1, 100));
    }
}

Output

[1, 2, 3, 4]
[1, 1, 2, 3]
[1, 3]
[5]

Time: O(n log n). Space: O(n).

Optimal: binary search for the left edge of the window

Because the array is sorted, the answer is always a contiguous window of length k, so only its starting index left is needed, which lies between 0 and n - k inclusive. Binary search over that range. Compare the two elements that change places when the window slides one step right: arr[mid] (the window's first element) and arr[mid + k] (the first element just outside). If x - arr[mid] > arr[mid + k] - x, the left end is strictly farther from x, so the window should move right: lo = mid + 1. Otherwise (including a tie, which favors the smaller value) the window can stay or move left: hi = mid.

Java

import java.util.ArrayList;
import java.util.List;

class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        int n = arr.length;
        k = Math.min(k, n);
        int lo = 0, hi = n - k;
        while (lo < hi) {
            int mid = lo + (hi - lo) / 2;
            if ((long) x - arr[mid] > (long) arr[mid + k] - x) {
                lo = mid + 1;
            } else {
                hi = mid;
            }
        }
        List<Integer> res = new ArrayList<>(k);
        for (int i = lo; i < lo + k; i++) {
            res.add(arr[i]);
        }
        return res;
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.findClosestElements(new int[] {1, 2, 3, 4, 5}, 4, 3));
        System.out.println(s.findClosestElements(new int[] {1, 1, 2, 3, 4, 5}, 4, -1));
        System.out.println(s.findClosestElements(new int[] {1, 3}, 2, 2));
        System.out.println(s.findClosestElements(new int[] {5}, 1, 100));
    }
}

Output

[1, 2, 3, 4]
[1, 1, 2, 3]
[1, 3]
[5]

Time: O(log(n - k) + k): the search, then copying the k results. Space: O(1) besides the output.

Since mid < hi <= n - k, mid + k <= n - 1, so the comparison never reads out of bounds. Both differences are computed in long so extreme values cannot overflow int. LeetCode's Java signature returns a List<Integer>, which prints as [1, 2, 3, 4].