LeetCode 658: Find K Closest Elements

Problem Restatement

We are given:

We must return the k elements closest to x.

The returned elements must remain sorted in ascending order.

If two numbers are equally close, we prefer the smaller number.

The problem defines closeness using:

|a - x| < |b - x|

or, when distances are equal:

a < b

Input and Output

Example function shape:

def findClosestElements(arr: list[int], k: int, x: int) -> list[int]:
    ...

Examples

Consider:

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

The closest elements to 3 are:

[1, 2, 3, 4]

Distances:

The four closest elements are:

[1, 2, 3, 4]

Another example:

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

Since x is left of all elements, the closest values are the first four elements:

[1, 2, 3, 4]

Another example:

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

The answer is:

[1, 2, 3, 4]

Even though both 1 and 5 are distance 2 away, we prefer the smaller value 1.

First Thought: Sort by Distance

A direct approach is:

1. Compute the distance of every element to x.
2. Sort by:
   • Distance
   • Value itself for tie-breaking
3. Take the first k elements.
4. Sort them again in ascending order.

This works, but it ignores an important fact:

arr is already sorted

We should use that structure.

Key Insight

The final answer is always a continuous subarray.

Suppose we choose k closest elements.

Because the array is sorted, if two elements belong to the answer, then every element between them must also belong to the answer.

So instead of choosing arbitrary elements, we only need to find:

the best window of length k

If the array length is n, then the starting index can range from:

0

to:

n - k

We can binary search for the correct starting position.

Comparing Two Windows

Suppose we are deciding between two neighboring windows:

arr[left : left + k]

and:

arr[left + 1 : left + 1 + k]

The only difference is:

So we compare:

x - arr[left]

and:

arr[left + k] - x

If the left side is larger, then the right endpoint is closer, so we should move the window right.

Otherwise, keep the current window.

Algorithm

Use binary search on the starting index of the window.

Initialize:

left = 0
right = len(arr) - k

While:

left < right

compute:

mid = (left + right) // 2

Then compare:

x - arr[mid]

and:

arr[mid + k] - x

If the left side is larger:

left = mid + 1

Otherwise:

right = mid

At the end, left is the best window start.

Return:

arr[left:left + k]

Correctness

The algorithm searches over all possible windows of length k.

For any neighboring windows:

arr[i : i+k]

and:

arr[i+1 : i+1+k]

the only difference is that the first contains arr[i], while the second contains arr[i+k].

If:

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

then arr[i+k] is closer to x than arr[i], so the better window must start to the right.

Otherwise, the current window is at least as good as the right neighbor. In the tie case, the smaller element should be preferred, so the left window remains better.

This comparison allows binary search because once the optimal window lies to the right, all earlier windows are worse.

The binary search eventually finds the unique smallest starting index of the optimal window. Returning the subarray of length k starting there gives exactly the k closest elements in sorted order.

Complexity

The returned slice itself contains k elements.

Implementation

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]

Code Explanation

The search space is the possible starting positions of a window of size k.

left = 0
right = len(arr) - k

If the array has length n, then the last valid window starts at:

n - k

We binary search this range.

while left < right:

At each step:

mid = (left + right) // 2

We compare the left endpoint removed from the current window with the new right endpoint added by shifting the window right.

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

If this condition is true, the right endpoint is closer, so the better window starts later.

left = mid + 1

Otherwise, the current window is still optimal or tied in a favorable way.

right = mid

Finally:

return arr[left:left + k]

returns the optimal window.

Testing

def run_tests():
    s = Solution()

    assert s.findClosestElements(
        [1, 2, 3, 4, 5],
        4,
        3,
    ) == [1, 2, 3, 4]

    assert s.findClosestElements(
        [1, 2, 3, 4, 5],
        4,
        -1,
    ) == [1, 2, 3, 4]

    assert s.findClosestElements(
        [1, 2, 3, 4, 5],
        4,
        10,
    ) == [2, 3, 4, 5]

    assert s.findClosestElements(
        [1, 1, 2, 3, 4, 5],
        4,
        3,
    ) == [1, 2, 3, 4]

    assert s.findClosestElements(
        [1],
        1,
        0,
    ) == [1]

    assert s.findClosestElements(
        [1, 2, 3, 4],
        2,
        2,
    ) == [1, 2]

    print("all tests passed")

run_tests()

Test meaning: