658. Find K Closest Elements


Problem Description

Given a sorted array of integers arr, and two integers k and x, the task is to find the k closest integers to x in the array. The result should be returned in ascending order. To determine which integers are the closest, we follow these rules:

  1. An integer a is considered closer to x than an integer b if the absolute difference |a - x| is less than |b - x|.
  2. If the absolute differences are the same, then a is closer to x than b if a is smaller than b.

In essence, the problem asks us to find a subsequence of the array that contains integers that are nearest to x, with a special emphasis on the absolute difference and the value of the integers themselves when the differences are the same.

Intuition

To solve the problem, we need an efficient way to find the subsequence of length k out of the sorted array that is closest to the target integer x.

Method 1: Sort A brute force approach would be to sort the elements of the array based on their distance from x. After sorting, we can simply take the first k elements. However, this method is not optimal in terms of time complexity because sorting would take O(n log n) time.

Method 2: Binary search A more efficient approach utilizes the fact that the array is already sorted. We know that the k closest integers we are looking for must form a consecutive subsequence in the array. Instead of looking for each integer separately, we should look for the starting index of this subsequence. We can use a binary search strategy to find this starting index. The high-level steps are:

  1. Set two pointers: left at the start of the array and right at the position len(arr) - k because any subsequence starting beyond this point does not have enough elements to be k long.
  2. Perform a binary search between left and right:
    • Calculate the middle index mid.
    • Check if x is closer to arr[mid] or to arr[mid + k].
    • If x is closer to arr[mid] or at the same distance to both, we discard the elements to the right of mid + k because the best starting point for our subsequence must be to the left or inclusive of mid.
    • If x is closer to arr[mid + k], we discard the elements to the left of mid since the best starting point must be to the right.
  3. Keep narrowing down until left is equal to right, indicating that we have found the starting index of the k closest elements.

Learn more about Two Pointers, Binary Search, Sorting, Sliding Window and Heap (Priority Queue) patterns.

Not Sure What to Study? Take the 2-min Quiz to Find Your Missing Piece:

Which algorithm should you use to find a node that is close to the root of the tree?

Solution Approach

The solution implements a binary search algorithm to efficiently find the starting index for the k closest elements to x. Here's a step-by-step breakdown of the approach and the code implementation:

  1. We begin by initializing two pointers, left at 0 and right at len(arr) - k. This is because the k elements we seek must form consecutive entrants in the array, and starting any further to the right would leave insufficient elements to reach a count of k.

  2. We enter a loop that continues until left is no longer less than right, indicating that we have narrowed down to the exact starting index for the k closest elements.

  3. Inside the loop, we find the middle index between left and right using the formula mid = (left + right) >> 1. The >> 1 is a bitwise right shift that effectively divides the sum by 2.

  4. With mid established, we check the distance between x and the current middle value at arr[mid] compared to the distance between x and the potential upper bound for our sequence at arr[mid + k].

    • If the element at mid is closer to x (or the elements are equally distant), we know that the k closest elements can't start any index higher than mid. Thus we set right to mid, moving our search space leftward.
    • Conversely, if the element at mid + k is closer to x, it means the sequence starts further to the right, so we move our left pointer up to mid + 1.
  5. By repeatedly halving our search space, we eventually arrive at a point where left equals right. This index represents the start of the subarray containing the k closest numbers to x.

  6. Once we have the starting index (left), we slice the original array from left to left + k to select the k closest elements.

  7. This slice is then returned, and since the original array was sorted, this resulting slice is guaranteed to be sorted as well, yielding the correct answer.

Using this binary search approach, we achieve a time complexity of O(log(n - k)) for finding the subsequence, which is an optimization over a complete sort-based method that would require O(n log n) time. This optimization is significant for large datasets where the difference in time complexity could be substantial.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

How many times is a tree node visited in a depth first search?

Example Walkthrough

Let us walk through a simple example to illustrate the solution approach.

Example

Assume we have an array arr = [1, 3, 4, 7, 8, 9], and we want to find k = 3 closest integers to x = 5.

Initial Setup

  • We begin by initializing two pointers, left is set to 0 and right is set to len(arr) - k, which in this case is 6 - 3 = 3. The subarray [7, 8, 9] is the last possible sequence of k numbers, and we will not consider starting indices beyond this.

Binary Search

  • The binary search begins. We calculate the middle index between left and right. The initial mid is (0 + 3) / 2 = 1.5, rounded down to 1.
  • We then compare the distance of x = 5 from arr[mid] = 3 and from arr[mid + k] = arr[4] = 8. The distance from 3 is |5 - 3| = 2 and from 8 is |5 - 8| = 3.
  • Since 3 is closer to 5 than 8 is, we set right to mid, which is now 1. Our search space for starting indices is now [0, 1].

Narrowing Down

  • We perform another iteration of the binary search. Our new mid is (0 + 1) / 2 = 0.5, rounded down to 0.
  • Again, we compare distances from x = 5: arr[mid] = 1 has a distance of 4, and arr[mid + k] = arr[3] = 7 has a distance of 2.
  • Since 7 is closer to 5 than 1 is, we move the left pointer to mid + 1, which is 1.

Convergence

  • Because left now equals right, we have converged to the starting index for the k closest elements. Our starting index is 1.

Obtaining the Result

  • We slice the original array from left to left + k, which yields [3, 4, 7].

Conclusion

From our example using the array arr = [1, 3, 4, 7, 8, 9] and the values k = 3 and x = 5, we have walked through the binary search approach and determined that the k closest integers to x in this array are [3, 4, 7], as illustrated by the binary search logic to find the starting index and the subsequent slice of the array. This method efficiently finds the correct subsequence without fully sorting the array based on proximity to x.

Not Sure What to Study? Take the 2-min Quiz:

How does merge sort divide the problem into subproblems?

Python Solution

1from typing import List
2
3class Solution:
4    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
5        # Initialize the left and right pointers for binary search.
6        # The right pointer is set to the highest starting index for the sliding window.
7        left, right = 0, len(arr) - k
8
9        # Perform binary search to find the left bound of the k closest elements.
10        while left < right:
11            # Calculate the middle index between left and right.
12            mid = (left + right) // 2
13          
14            # Check the distance from the x to the middle element and the element at mid + k position.
15            # If the element at mid is closer to x or equal in distance compared to the element at mid + k,
16            # we move the right pointer to mid. Otherwise, we adjust the left pointer to mid + 1.
17            if x - arr[mid] <= arr[mid + k] - x:
18                right = mid
19            else:
20                left = mid + 1
21      
22        # Extract the subarray from the left index of size k, which will be the k closest elements.
23        return arr[left:left + k]
24
25# Example usage:
26# sol = Solution()
27# result = sol.findClosestElements([1, 2, 3, 4, 5], k=4, x=3)
28# print(result)  # Output should be [1, 2, 3, 4]
29

Java Solution

1import java.util.List;
2import java.util.ArrayList;
3
4class Solution {
5    public List<Integer> findClosestElements(int[] arr, int k, int x) {
6        // Initialize the left and right pointers for binary search.
7        int left = 0;
8        int right = arr.length - k;
9
10        // Continue searching until the search space is reduced to a single element.
11        while (left < right) {
12            // Calculate the middle index.
13            int mid = left + (right - left) / 2; // Using left + (right - left) / 2 to avoid potential overflow.
14
15            // If the distance to the left is less than or equal to the distance to the right,
16            // we need to move towards the left (lower indices).
17            if (x - arr[mid] <= arr[mid + k] - x) {
18                right = mid;
19            } else {
20                // Otherwise, move right (higher indices).
21                left = mid + 1;
22            }
23        }
24
25        // Create a list to store the k closest elements.
26        List<Integer> result = new ArrayList<>();
27        // Add the closest elements to the list, starting from the left pointer.
28        for (int i = left; i < left + k; ++i) {
29            result.add(arr[i]);
30        }
31
32        // Return the list of k closest elements.
33        return result;
34    }
35}
36

C++ Solution

1#include <vector>
2
3class Solution {
4public:
5    std::vector<int> findClosestElements(std::vector<int>& arr, int k, int x) {
6        // Initialize the binary search bounds
7        int left = 0;
8        int right = arr.size() - k;
9
10        // Perform binary search to find the start index of the k closest elements
11        while (left < right) {
12            // Calculate mid index (avoid potential overflow by using left + (right-left)/2)
13            int mid = left + (right - left) / 2;
14          
15            // Compare the differences between x and elements at mid index and mid+k index
16            // The goal is to find the smallest window such that the elements are closest to x
17            if (x - arr[mid] <= arr[mid + k] - x) {
18                // If the element at mid index is closer to x, or equally close
19                // as the element at mid+k index (prefer the smaller element),
20                // move the right bound to mid
21                right = mid;
22            } else {
23                // Otherwise, if the element at mid+k index is closer to x,
24                // move the left bound to mid + 1
25                left = mid + 1;
26            }
27        }
28
29        // Create and return a vector of k closest elements starting from the 'left' index
30        return std::vector<int>(arr.begin() + left, arr.begin() + left + k);
31    }
32};
33

Typescript Solution

1/**
2 * Finds the k closest elements to a given target x in a sorted array.
3 * 
4 * @param {number[]} arr - The sorted array of numbers.
5 * @param {number} k - The number of closest elements to find.
6 * @param {number} x - The target number to find the closest elements to.
7 * @returns {number[]} - An array of k closest elements to the target.
8 */
9function findClosestElements(arr: number[], k: number, x: number): number[] {
10    // Initialize two pointers for binary search.
11    let leftPointer = 0;
12    let rightPointer = arr.length - k;
13
14    // Binary search to find the start index of the k closest elements.
15    while (leftPointer < rightPointer) {
16        const midPointer = (leftPointer + rightPointer) >> 1; // Equivalent to Math.floor((left + right) / 2)
17  
18        // If the element at the middle is closer to or as close as the element k away from it,
19        // then the closer elements are to the left side of the array.
20        // Otherwise, they are to the right, and we move the left pointer one step past the middle.
21        if (x - arr[midPointer] <= arr[midPointer + k] - x) {
22            rightPointer = midPointer;
23        } else {
24            leftPointer = midPointer + 1;
25        }
26    }
27
28    // Slice the array from the left pointer to get the k closest elements.
29    return arr.slice(leftPointer, leftPointer + k);
30}
31
Fast Track Your Learning with Our Quick Skills Quiz:

Which of these pictures shows the visit order of a depth-first search?

Time and Space Complexity

The provided algorithm is a binary search approach for finding the k closest elements to x in the sorted array arr. Here's the analysis:

  • Time Complexity: The time complexity of this algorithm is O(log(N - k)) where N is the number of elements in the array. This is because the binary search is performed within a range that's reduced by k elements (since we're always considering k contiguous elements as a potential answer).

    The algorithm performs a binary search by repeatedly halving the search space, which is initially N - k. It does not iterate through all N elements, but narrows down the search to the correct starting point for the sequence of k closest elements. After finding the starting point, it returns the subarray in constant time, as array slicing in Python is done in O(1) time.

    Hence, the iterative process of the binary search dominates the running time, which leads to the O(log(N - k)) time complexity.

  • Space Complexity: The space complexity of this algorithm is O(1). The code uses a fixed amount of extra space for variables left, right, and mid. The returned result does not count towards space complexity as it is part of the output.

Please note that array slicing in the return statement does not create a deep copy of the elements but instead returns a reference to the same elements in the original array, thus no extra space proportional to k or N is used in this operation.

Learn more about how to find time and space complexity quickly.


Recommended Readings


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫