658.find-k-closest-elements.md

658. Find K Closest Elements

Clarification Questions

• k range?
• Does x always exist in the array?

Test Cases

Normal Cases

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

Edge / Corner Cases

• x does not exist in the array.

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

• There are some number with the same distance between x.

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

Pitfalls

The problem is asking k closest integers to x, however, x may not be in the array, so we can't just find the index of x and expand the range to find the k closest integers.

   [y y X y y y y y y y]
|------ k ------|

    X  [y y y y y y y y y y]
|-- k --|

Heap

Just sort by the distance of x, and find the k closest (k smallest) elements.

fun findClosestElements(arr: IntArray, k: Int, x: Int): List<Int> {
    // To find `k` smallest elements, we can use a max heap of fixed size `k`.
    val maxHeap = PriorityQueue<Int>() { n1, n2 -> 
        val d1 = abs(x - n1)
        val d2 = abs(x - n2)

        if (d1 == d2) n2 - n1
        else d2 - d1
    }
    for (n in arr) {
        maxHeap.add(n)
        if (maxHeap.size > k) {
            maxHeap.poll()
        }
    }
    val results = mutableListOf<Int>()
    while (maxHeap.isNotEmpty()) {
        results.add(maxHeap.poll())
    }
    results.sort()
    return results
}

• Time Complexity: O(n * lg k + k * lg k).
• Space Complexity: O(k).

Two Pointers (Binary Search)

Idea!! We're looking for the window size is k and we initialize the window size as array size, then we try to shrink the window size to k.

We set two points at the first and last position, and check the distance of x with the two points, then shrink the range of longer distance until the size is k. If the distance is equal, then we shrink the right pointer because of this rule: |a - x| == |b - x| and a < b in the problem description.

fun findClosestElements(arr: IntArray, k: Int, x: Int): List<Int> {
    var left = 0
    var right = arr.size - 1
    while (right - left + 1 > k) { // Equivalent to `while ((right - left) >= k)`
        val leftDiff = abs(arr[left] - x)
        val rightDiff = abs(arr[right] - x)
        if (leftDiff > rightDiff) left++
        else right--
    }
    return arr.slice(left..right)
}

/**
 Why `right - left >= k` works? It will stop when `right - left == k - 1` which is equivalent to `right - left + 1 == k`.
 */

• Time Complexity: O(n).
• Space Complexity: O(k) for results.

Test Cases

[1, 2, 5, 7, 10] k = 1, 2, 4; x = -1, 4, 6