找到 K 个最接近的元素

标签: 数组 双指针 二分查找 排序 滑动窗口 堆(优先队列)

难度: Medium

给定一个 排序好 的数组 arr ,两个整数 kx ,从数组中找到最靠近 x(两数之差最小)的 k 个数。返回的结果必须要是按升序排好的。

整数 a 比整数 b 更接近 x 需要满足:

  • |a - x| < |b - x| 或者
  • |a - x| == |b - x|a < b

示例 1:

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

示例 2:

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

提示:

  • 1 <= k <= arr.length
  • 1 <= arr.length <= 104
  • arr 按 升序 排列
  • -104 <= arr[i], x <= 104

Submission

运行时间: 25 ms

内存: 17.0 MB

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 - left) // 2
            if x - arr[mid] > arr[mid + k] - x:
                left = mid + 1
            else:
                right = mid
        
        return arr[left:left+k]

Explain

这个题解使用了二分查找的思路。我们要在数组中找到一个位置 left,使得区间 [left, left+k) 包含 k 个最接近 x 的元素。由于数组已经按升序排列,所以只需要比较 x 与区间两端点的差值大小即可。如果 x 与左端点的差值更大,说明 left 需要右移;否则说明 left 已经是最优区间的左端点,可以结束查找。

时间复杂度: O(log n)

空间复杂度: O(1)

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 - left) // 2
            # 如果 x 与左端点的差值更大,说明左端点需要右移
            if x - arr[mid] > arr[mid + k] - x:
                left = mid + 1
            # 否则说明 left 已经是最优区间的左端点,可以结束查找
            else:
                right = mid
        
        # 返回最接近 x 的 k 个元素组成的区间
        return arr[left:left+k]

Explore

这个条件用来比较x与区间左端点`arr[mid]`和区间右端点`arr[mid+k]`的距离。如果`x - arr[mid]`大于`arr[mid+k] - x`,意味着x更接近右端点,因此需要将搜索区间向右移动(即移动左端点`left`),以包含更接近x的元素。反之,如果x更接近左端点或与两端点的距离相等,维持当前的左端点位置,因为它可能就是最终的最优区间左端点。

当x非常接近`arr[mid]`或`arr[mid+k]`时,`x - arr[mid]`与`arr[mid + k] - x`的比较仍然有效。如果x几乎等于`arr[mid]`,这意味着左端点可能已经非常接近最优位置,因为x与左端点的距离小于或等于与右端点的距离。此时算法可能会决定维持当前的左端点。如果x接近`arr[mid+k]`,则右端点的距离可能更小,导致算法调整左端点向右,直到找到最接近x的整体区间。

在二分查找中,`left`和`right`变量定义了搜索区间的边界。随着算法的进行,`left`和`right`逐渐逼近最优区间的左端点。最终,`left`变量停留在使区间`[left, left+k)`包含最接近x的k个元素的位置。由于我们总是调整`left`或`right`以逼近这一最优位置,所以当`left`停止移动时,它指向的就是这一区间的开始。`right`通常作为搜索的上界,其值不一定是最优区间的一部分。

如果`x`远离`arr`的所有元素,二分查找仍然高效。在这种情况下,二分查找将快速收敛到数组的一个端点。例如,如果x远小于`arr`中的最小元素,二分查找将快速定位到数组的最左端;如果x远大于数组中的最大元素,则查找将快速定位到数组的右端(考虑到要取k个元素,实际上是`len(arr)-k`)。在这两种极端情况下,二分查找的步骤数仍然是对数级别的,确保了算法的效率。