Want Help Cracking FAANG?

(Then click this)

×
Back to Question Bank

658. Find K Closest Elements - Leetcode Solution

Code Implementation

from typing import List

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        left, right = 0, len(arr) - k
        while left < right:
            mid = (left + right) // 2
            # Compare distances between x and the edges of the window
            if x - arr[mid] > arr[mid + k] - x:
                left = mid + 1
            else:
                right = mid
        return arr[left:left + k]
      
#include <vector>
#include <algorithm>
using namespace std;

class Solution {
public:
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        int left = 0, right = arr.size() - k;
        while (left < right) {
            int mid = (left + right) / 2;
            if (x - arr[mid] > arr[mid + k] - x)
                left = mid + 1;
            else
                right = mid;
        }
        return vector<int>(arr.begin() + left, arr.begin() + left + k);
    }
};
      
import java.util.*;

class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        int left = 0, right = arr.length - k;
        while (left < right) {
            int mid = (left + right) / 2;
            if (x - arr[mid] > arr[mid + k] - x) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }
        List<Integer> result = new ArrayList<>();
        for (int i = left; i < left + k; i++) {
            result.add(arr[i]);
        }
        return result;
    }
}
      
/**
 * @param {number[]} arr
 * @param {number} k
 * @param {number} x
 * @return {number[]}
 */
var findClosestElements = function(arr, k, x) {
    let left = 0, right = arr.length - k;
    while (left < right) {
        let 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);
};
      

Problem Description

You are given a sorted array of integers arr, an integer k, and an integer x. Your task is to find the k elements in arr that are closest to x. The result should be sorted in ascending order.

  • If there is a tie (two numbers are equally close to x), the smaller element should be preferred.
  • The input array arr is already sorted in ascending order.
  • Return the answer as a list of integers.
  • All elements are unique, and you should not reuse elements.
  • There is always exactly one valid solution.

Thought Process

To solve this problem, the first idea that comes to mind is to check the distance of every element in arr from x, sort them by distance, and pick the closest k elements. This brute-force method is straightforward but inefficient, especially for large arrays.

However, since the array is already sorted, we should look for a more optimal solution. Instead of comparing all elements, we can leverage the sorted property to efficiently find the window of size k that contains the closest elements to x. This leads us to consider binary search for efficiency.

The challenge is to identify the correct starting point for the window of k elements. We need a method that efficiently narrows down the options without checking all possible windows.

Solution Approach

We use a binary search strategy to find the starting index of the window of length k that contains the closest elements to x.

  1. Define the Search Space: The possible starting indices for a window of length k are from 0 to len(arr) - k.
  2. Binary Search: For each possible window starting at index mid, compare the distances between x and the left edge (arr[mid]) and the right edge (arr[mid + k]).
    • If x - arr[mid] > arr[mid + k] - x, then x is closer to the right edge, so move the window to the right (left = mid + 1).
    • Otherwise, move the window to the left (right = mid).
  3. Termination: When left == right, we've found the optimal window. The answer is arr[left:left + k].

This approach works efficiently because the sorted property of arr ensures that the closest window will be contiguous and can be found using binary search.

Example Walkthrough

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

  1. Initial search space: left = 0, right = 1 (since len(arr) - k = 1)
  2. First iteration:
    • mid = 0
    • Compare x - arr[mid] = 3 - 1 = 2 and arr[mid + k] - x = 5 - 3 = 2
    • Since 2 > 2 is false, move right = mid = 0
  3. End condition: left == right == 0
  4. Result: The window is arr[0:4] = [1, 2, 3, 4]
  5. Why? These are the 4 numbers closest to 3. If there is a tie (as between 1 and 5), the smaller number is chosen.

Time and Space Complexity

  • Brute-force approach:
    • Time: O(n log n) (for sorting by distance)
    • Space: O(n) (for storing distances and sorting)
  • Optimized binary search approach:
    • Time: O(log(n - k) + k). Binary search takes O(log(n - k)) and copying the window takes O(k).
    • Space: O(1) extra (excluding the output)

The optimized approach is much more efficient, especially for large arrays.

Summary

By leveraging the sorted property of the input array, we avoid brute-force methods and use binary search to efficiently find the window of k closest elements to x. This approach is elegant because it reduces the problem to a simple window search, making the solution both fast and easy to understand.