LeetCode 658 - Find K Closest Elements

The problem gives us a sorted integer array arr, along with two integers, k and x. We need to return exactly k elements from the array that are closest to x.

leetcodemediumarraytwo-pointersbinary-searchsliding-windowsortingheap-(priority-queue)

LeetCode Problem 658

Difficulty: 🟡 Medium
Topics: Array, Two Pointers, Binary Search, Sliding Window, Sorting, Heap (Priority Queue)

Solution

Problem Understanding

The problem gives us a sorted integer array arr, along with two integers, k and x. We need to return exactly k elements from the array that are closest to x.

The definition of "closest" is important:

  • An element a is closer than b if |a - x| < |b - x|
  • If both elements are equally far away from x, then the smaller element is considered closer

The returned result must also remain sorted in ascending order.

Since the input array is already sorted, this is a major clue that we should avoid unnecessary sorting operations and instead take advantage of the ordering structure.

For example, if:

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

then the four closest values to 3 are:

[1,2,3,4]

Even though 5 is also near 3, it is farther away than 1.

The constraints are:

  • 1 <= k <= arr.length
  • arr.length <= 10^4
  • arr is already sorted

An array size of 10^4 is not extremely large, so even an O(n log n) solution would pass comfortably. However, because the array is sorted, we can do better with binary search.

Several edge cases are important:

  • x may be smaller than every element in the array
  • x may be larger than every element
  • Multiple elements may have the same distance from x
  • k may equal the entire array length
  • Duplicate values may exist

These cases can easily break naive implementations if tie handling is incorrect or if boundaries are not managed carefully.

Approaches

Brute Force Approach

A straightforward solution is to compute the distance of every element from x, then sort the array according to:

  1. Absolute distance from x
  2. Element value itself for tie breaking

After sorting, we take the first k elements and finally sort those selected elements again to satisfy the output ordering requirement.

This approach works because it directly follows the problem definition. Every element receives a ranking based on closeness to x.

However, this solution performs unnecessary sorting work. Even though the array is already sorted, we ignore that structure and reorder everything from scratch.

Key Insight for the Optimal Solution

The important observation is that the answer must form a contiguous subarray of length k.

Because the array is sorted, if we pick two elements far apart while skipping a middle element, the skipped element would always be at least as close as one of the chosen endpoints. Therefore, the optimal answer is always a continuous window.

This changes the problem from:

"Which individual elements should we pick?"

to:

"Where should the window of size k start?"

If the window has size k, then its starting index can range from:

0 to len(arr) - k

We can binary search over these possible starting positions.

For a candidate window starting at index mid, we compare:

arr[mid]

and

arr[mid + k]

These represent the elements just outside competing windows.

If:

x - arr[mid] > arr[mid + k] - x

then the right side is closer, so we should move the window right.

Otherwise, we keep or move left.

This gives an elegant O(log(n-k) + k) solution.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(n log n) O(n) Sorts all elements by distance from x
Optimal O(log(n-k) + k) O(1) extra space Binary searches for the best window

Algorithm Walkthrough

Optimal Binary Search Window Algorithm

  1. Initialize two pointers, left = 0 and right = len(arr) - k.

We are searching for the starting index of the best window of size k. 2. Perform binary search while left < right.

At each step, compute:

mid = (left + right) // 2
  1. Compare the distances between x and the competing boundary elements.

The two relevant elements are:

arr[mid]
arr[mid + k]

These correspond to two overlapping windows:

  • Window starting at mid
  • Window starting at mid + 1
  1. Decide which direction is better.

If:

x - arr[mid] > arr[mid + k] - x

then the right boundary element is closer to x, so shifting the window right improves the result.

Move:

left = mid + 1
  1. Otherwise, keep the current window or search left.

Move:

right = mid
  1. Continue until left == right.

At this point, we have found the optimal starting index. 7. Return the subarray:

arr[left:left+k]

Why it works

The binary search works because the quality of the window changes monotonically. As we move the window from left to right, there is a transition point where windows stop improving and begin worsening.

By comparing arr[mid] and arr[mid + k], we determine which neighboring window is better. Since the array is sorted, this comparison is sufficient to eliminate half the search space each iteration.

The final window is guaranteed to contain the k closest elements because every discarded window was proven worse during binary search.

Python Solution

from typing import List

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        left = 0
        right = len(arr) - k

        while left < right:
            mid = (left + right) // 2

            if x - arr[mid] > arr[mid + k] - x:
                left = mid + 1
            else:
                right = mid

        return arr[left:left + k]

The implementation directly follows the binary search window strategy.

The variables left and right define the search space for possible window starting indices. Since the window length is fixed at k, the largest valid starting index is len(arr) - k.

Inside the loop, mid represents a candidate starting position.

The key comparison:

x - arr[mid] > arr[mid + k] - x

checks whether the element just beyond the right side of the current window is closer to x than the leftmost element currently inside the window.

If the right candidate is better, the optimal window must start farther right, so we move left.

Otherwise, the current window is at least as good as the shifted one, so we keep searching left.

When the loop terminates, left stores the optimal starting index. We then return the corresponding slice of length k.

Go Solution

func findClosestElements(arr []int, k int, x int) []int {
	left := 0
	right := len(arr) - k

	for left < right {
		mid := (left + right) / 2

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

	return arr[left : left+k]
}

The Go implementation mirrors the Python logic very closely.

Go slices make extracting the final subarray efficient because slicing does not copy elements immediately. The returned slice references the appropriate section of the original array.

Unlike Python, Go does not require importing typing utilities. Integer arithmetic is also straightforward because the constraints are small enough to avoid overflow concerns.

Worked Examples

Example 1

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

We search for the best window of size 4.

Possible windows:

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

Binary Search Trace

left right mid arr[mid] arr[mid+k] Comparison Action
0 1 0 1 5 3-1 > 5-3 → 2 > 2 is false right = mid

Loop ends because:

left == right == 0

Return:

arr[0:4] = [1,2,3,4]

Example 2

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

Possible windows:

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

Binary Search Trace

left right mid arr[mid] arr[mid+k] Comparison Action
0 2 1 1 5 -1-1 > 5-(-1) → -2 > 6 is false right = 1
0 1 0 1 4 -1-1 > 4-(-1) → -2 > 5 is false right = 0

Loop ends:

left == right == 0

Return:

[1,1,2,3]

Complexity Analysis

Measure Complexity Explanation
Time O(log(n-k) + k) Binary search over window positions plus output construction
Space O(1) extra space Only pointer variables are used

The binary search examines at most log(n-k) candidate windows. After finding the optimal starting index, extracting the final window requires returning k elements.

No auxiliary data structures are required, so the extra space usage remains constant.

Test Cases

from typing import List

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        left = 0
        right = len(arr) - k

        while left < right:
            mid = (left + right) // 2

            if x - arr[mid] > arr[mid + k] - x:
                left = mid + 1
            else:
                right = mid

        return arr[left:left + k]

solution = Solution()

assert solution.findClosestElements([1,2,3,4,5], 4, 3) == [1,2,3,4]  # basic centered case

assert solution.findClosestElements([1,1,2,3,4,5], 4, -1) == [1,1,2,3]  # x smaller than all values

assert solution.findClosestElements([1,2,3,4,5], 4, 10) == [2,3,4,5]  # x larger than all values

assert solution.findClosestElements([1], 1, 0) == [1]  # single element array

assert solution.findClosestElements([1,2,3,4,5], 5, 3) == [1,2,3,4,5]  # k equals array size

assert solution.findClosestElements([1,2,3,4,5], 1, 3) == [3]  # exact match

assert solution.findClosestElements([1,2,3,4,5], 1, 6) == [5]  # choose largest element

assert solution.findClosestElements([1,2,3,4,5], 1, -2) == [1]  # choose smallest element

assert solution.findClosestElements([1,2,2,2,3,3,4], 3, 2) == [2,2,2]  # duplicates around target

assert solution.findClosestElements([1,3], 1, 2) == [1]  # tie broken toward smaller value

Test Summary

Test Why
[1,2,3,4,5], k=4, x=3 Standard balanced example
[1,1,2,3,4,5], k=4, x=-1 Target smaller than all elements
[1,2,3,4,5], k=4, x=10 Target larger than all elements
[1], k=1, x=0 Minimum input size
k == len(arr) Entire array must be returned
k == 1 with exact match Smallest window size
k == 1 with large target Right boundary handling
k == 1 with small target Left boundary handling
Duplicate values around target Ensures duplicate handling is correct
Tie case [1,3], x=2 Verifies smaller value wins ties

Edge Cases

One important edge case occurs when x is smaller than every element in the array. In this situation, the closest elements are simply the first k values. A buggy implementation might incorrectly move the window right because of negative arithmetic or improper comparisons. The binary search comparison naturally handles this case because all left-side elements remain closer than farther right elements.

Another critical edge case occurs when x is larger than every array element. Here, the answer should be the final window of size k. Incorrect boundary calculations can easily produce index errors near the end of the array. Using right = len(arr) - k guarantees every candidate window remains valid.

Tie situations are also especially important. Suppose x = 2 and the candidates are 1 and 3. Both are equally distant from 2, but the problem requires choosing the smaller element. The comparison:

if x - arr[mid] > arr[mid + k] - x

uses a strict greater-than operator rather than greater-than-or-equal. This subtle detail ensures ties favor the left side, which corresponds to smaller elements.

Another edge case involves duplicate values. If many identical values surround x, some implementations may incorrectly shrink or expand the window. Since this algorithm only compares window boundaries and never removes middle elements individually, duplicates are handled naturally and correctly.

Finally, when k equals the full array length, there is only one valid window. The algorithm immediately returns the entire array because the binary search range collapses to a single starting index of 0.