在问 · 算法题解

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

中等二分查找二分查找窗口左端点力扣原题 ↗

题目描述

输入一个已按升序排列的数组 arr,以及整数 k 和 x。

比较两个数 a、b 哪个更靠近 x:先比 |a - x| 与 |b - x|,较小的更近;二者相等时,数值较小的那个算更近。

按这个规则从 arr 中选出最靠近 x 的 k 个元素,并以升序返回。

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

解释:离 3 最近的是 3、2、4,再在 1 和 5 之间按平局规则选较小的 1。

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

解释:x 比所有元素都小,取最前面的 4 个。

数据范围

  • 1 <= k <= arr.length
  • 1 <= arr.length <= 10^4
  • arr 按升序排列。
  • -10^4 <= arr[i], x <= 10^4

解题思路:二分查找窗口左端点

答案一定是一段长度为 k 的连续窗口 arr[lo:lo+k],只需二分窗口左端点 lo,范围是 [0, n-k]。比较窗口左端 arr[mid] 和窗口外右边第一个数 arr[mid+k]:如果 x - arr[mid] > arr[mid+k] - x,说明右边那个数更近,窗口必须右移,lo = mid + 1;否则左端不差于右边,最优左端点在 mid 或更左。

为什么正确

最优窗口的左端点始终在 [lo, hi] 中。

复杂度分析

时间 O(log(n-k) + k),空间 O(1)(不计输出)。

易错点

比较时不能用绝对值 |x-arr[mid]| 与 |arr[mid+k]-x|:当 x 在窗口两端同一侧时绝对值会给出错误方向。用带符号的差即可,相等时偏向左边。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言,逻辑一致。

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        lo, hi = 0, len(arr) - k
        while lo < hi:
            mid = (lo + hi) // 2
            if x - arr[mid] > arr[mid + k] - x:
                lo = mid + 1
            else:
                hi = mid
        return arr[lo:lo + k]

载入编辑器,自己跑一遍 →

其他解法

按距离排序没有利用数组有序;双指针利用“答案是连续区间”从两端删元素;二分查找直接定位长为 k 的窗口的左端点,最快。

解法 2:双指针收缩窗口 时间 O(n),空间 O(1)

结果一定是一个长为 k 的连续区间。从整个数组 [lo, hi] 开始,每次删掉两端中离 x 更远的一个。距离相等时删右端,因为题目要求优先保留较小的数。区间长度缩到 k 时返回。

思路直观,没有二分的边界细节。代价是要删 n − k 次,k 很小时比二分慢。

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        lo, hi = 0, len(arr) - 1
        while hi - lo + 1 > k:
            if x - arr[lo] <= arr[hi] - x:
                hi -= 1
            else:
                lo += 1
        return arr[lo:hi + 1]

解法 3:按距离排序 时间 O(n log n),空间 O(n)

把元素按与 x 的距离从小到大排序,距离相同则较小的数在前。取前 k 个,再按数值升序排列后返回。

直接把题意翻译成排序规则,不依赖输入有序。代价是 O(n log n) 的排序和额外空间。

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        byDist = sorted(arr, key=lambda v: (abs(v - x), v))
        return sorted(byDist[:k])

二分查找 专题的其他题目

查看 二分查找 全部题目 →