658. Find K Closest Elements Medium

Given a sorted integer array arr, two integers k and x, return the k closest integers to x in the array. The result should also be sorted in ascending order.

An integer a is closer to x than an integer b if:

• |a - x| < |b - x|, or
• |a - x| == |b - x| and a < b

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

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

Approach

Input: A sorted array arr, and two integers k and x

Output: Return the subarray of k elements closest to x

This problem is a typical Binary Search + Sliding Window problem.

Although this problem seems like a sliding window type, since the array is sorted, we can use binary search to quickly find the left boundary of the window.

We can initialize the starting window range as [0, len(arr) - k].

Through binary search, we find mid and check if x - arr[mid] > arr[mid + k] - x.

If it's greater, it means the position mid is still relatively far from x, and the left boundary must be to the right of mid. So we move the left boundary to mid + 1.

If it's less than or equal, it means the position mid is closer to x, and the left boundary might be at mid or to the left of mid. We can continue to narrow the range to [left, mid] to find a better left boundary position.

Finally, left will be the starting point of the left boundary, and we directly return the subarray in the range [left, left + k].

Implementation

python

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        # Set the binary search range for the start of the window: from 0 to len(arr) - k
        # Because the window length is fixed at k, the start point can at most be len(arr) - k
        left, right = 0, len(arr) - k

        while left < right:
            mid = (left + right) // 2  # Current window start

            # Compare the distance of the two endpoints of the window to x to decide moving direction
            # If the left arr[mid] is further (further from x), move the window to the right
            # Otherwise, keep current mid or try to find a better one to the left
            if x - arr[mid] > arr[mid + k] - x:
                left = mid + 1
            else:
                right = mid

        # Finally, the start of the window is left, and the corresponding continuous k numbers are the answer
        return arr[left:left + k]

javascript

/**
 * @param {number[]} arr
 * @param {number} k
 * @param {number} x
 * @return {number[]}
 */
var findClosestElements = function(arr, k, x) {
    let left = 0;
    let right = arr.length - k;

    while (left < right) {
        const mid = Math.floor((left + right) / 2);

        if (x - arr[mid] > arr[mid + k] - x) {
            left = mid + 1;
        } else {
            right = mid;
        }
    }

    return arr.slice(left, left + k);
};

Complexity Analysis

• Time Complexity: O(log(n - k) + k)
• Space Complexity: O(1)

Links

658. Find K Closest Elements (English)

658. 找到 K 个最接近的元素 (Chinese)