[658][中等][二分] 找到 K 个最接近的元素

题目描述

658. 找到 K 个最接近的元素

给定一个排序好的数组，两个整数 k 和 x，从数组中找到最靠近 x（两数之差最小）的 k 个数。返回的结果必须要是按升序排好的。如果有两个数与 x 的差值一样，优先选择数值较小的那个数。

示例 1:

输入: [1,2,3,4,5], k=4, x=3
输出: [1,2,3,4]

示例 2:

输入: [1,2,3,4,5], k=4, x=-1
输出: [1,2,3,4]

说明:

• k 的值为正数，且总是小于给定排序数组的长度。

• 数组不为空，且长度不超过 104

• 数组里的每个元素与 x 的绝对值不超过 104

解题思路

首先用二分法定位到x的位置, x可能不在数组中, 需要对找到的位置的左右判别, 得到差距最小的位置. 然后以这个位置为中心, 使用双指针左右搜索.

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

        # 二分法确定`x`在数组中的位置
        left, right = 0, n - 1
        while left <= right:
            mid = (left + right) // 2
            if arr[mid] < x:
                left = mid + 1
            else:
                right = mid - 1

        # `left`为插入点， 判断`left`和`left - 1`两个位置哪个与`x`更相近
        # 以更相近的位置为中心位置, 向两侧扩展
        center = (left - 1) if left == n or (left > 0 and (x - arr[left - 1]) <= (arr[left] - x)) else left
        left = right = center
        count = 1
        while count < k:
            if left == 0:
                right += 1
                count += 1
            elif right == n - 1:
                left -= 1
                count += 1
            else:
                if (x - arr[left - 1]) <= (arr[right + 1] - x):
                    left -= 1
                else:
                    right += 1
                count += 1

        return arr[left: right + 1]