658 找到 K 个最接近的元素（大根堆、二分查找）

灰太狼 2022-08-29 10:08 104阅读 0赞

**1. 问题描述：**

给定一个排序好的数组 arr ，两个整数 k 和 x ，从数组中找到最靠近 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 <= 10 ^ 4  
数组里的每个元素与 x 的绝对值不超过 10 ^ 4  
来源：力扣（LeetCode）  
链接：https://leetcode-cn.com/problems/find-k-closest-elements

**2. 思路分析：**

① 分析题目可以知道我们需要求解差值前k小的元素，**求解前k小的元素**是一个经典问题，我们需要一个什么样的数据结构来支持插入和弹出当前与x差值最大的元素？可以想到我们可以使用堆来维护前k小/大的元素，所以对于这道题目来说可以使用**大根堆**来维护元素与x的差值，当堆顶元素个数大于k之后那么弹出堆顶中差值最大的元素，循环结束之后堆中就只剩下k个与x差值最小的元素，因为使用的是python语言所以我们需要在元素前面取一个负号表示大根堆（python中只有小根堆），堆中的元素需要维护两个信息，第一个是元素与x的差的绝对值，第二个是当前的元素值，最后将堆中第二个信息加入到答案即可。反正一看到需要维护前k小的元素那么立马想到堆来维护。

② 除了使用堆来维护前k小的思路之外，因为arr数组是有序的，所以还可以使用二分查找，我们可以通过二分先找到最靠近x的位置，这里可以**二分**大于等于x的位置r，并且我们需要先判断一下当前位置r与前一个位置r - 1哪一个位置更好那么就使用哪一个，因为有可能二分出来的这个位置不是最优的位置，所以我们需要先比较一下，然后使用**双指针**找到以当前分隔点r最靠近x的剩余k - 1个元素即可。

**3. 代码如下：**

**大根堆：**

from typing import List
    import heapq
    
    
    class Solution:
        # 使用大根堆来维护差值前k大的数
        def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
            q = list()
            for v in arr:
                # 这里需要特判一下当堆中元素个数等于k而且差值相等的时候跳过当前元素即可
                if len(q) == k and abs(v - x) == -q[0][0]: continue
                heapq.heappush(q, (-abs(v - x), v))
                if len(q) > k: heapq.heappop(q)
            res = list()
            # 当前列表存储的就是差值最小的k个数
            while q:
                t = q.pop()
                res.append(t[1])
            res.sort()
            return res

**二分查找：**

from typing import List
    
    
    class Solution:
        # 时间复杂度为O(logn + k)
        def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
            l, r = 0, len(arr) - 1
            while l < r:
                mid = l + r >> 1
                if arr[mid] >= x:
                    r = mid
                else:
                    l = mid + 1
            # 有可能左边的值更优秀所以需要先判断一下
            if r:
                a, b = abs(arr[r - 1] - x), abs(arr[r] - x)
                if a <= b: r -= 1
            i, j = r, r
            for u in range(k - 1):
                if i - 1 < 0:
                    j += 1
                elif j + 1 == len(arr):
                    i -= 1
                else:
                    # 比较两个数字哪一个更好
                    a, b = abs(x - arr[i - 1]), abs(x - arr[j + 1])
                    if a <= b:
                        i -= 1
                    else:
                        j += 1
            res = list()
            for u in range(i, j + 1):
                res.append(arr[u])
            return res