[Leetcode] 658. Find K Closest Elements (二分法)

最新推荐文章于 2022-01-19 15:11:24 发布

Haskei

最新推荐文章于 2022-01-19 15:11:24 发布

阅读量155

点赞数

分类专栏：二分法

本文链接：https://blog.csdn.net/Haskei/article/details/107725075

版权

二分法专栏收录该内容

10 篇文章 0 订阅

订阅专栏

658. Find K Closest Elements（题目链接）

Medium

Given a sorted array arr, two integers k and x, find the k closest elements to x in the array. The result should also be sorted in ascending order. If there is a tie, the smaller elements are always preferred.

Example 1:

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

Example 2:

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

Constraints:

1 <= k <= arr.length
1 <= arr.length <= 10^4
Absolute value of elements in the array and x will not exceed 104

方法1：

用一般的二分法，找到不大于x的那个数，然后从这个数往左往右找k个最接近的数

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        if k <= 0:
            return []
        l, r = 0, len(arr)
        while l < r:
            mid = (l + r) // 2
            if arr[mid] < x:
                l = mid + 1
            else:
                r = mid
        l -= 1
        while k:
            if l >= 0 and r < len(arr):
                left = x - arr[l]
                right = arr[r] - x
                if left <= right:
                    l -= 1
                else:
                    r += 1
            elif l >= 0:
                l -= 1
            else:
                r += 1
            k -= 1
        return arr[l + 1:r]

方法2：

Binary search, time O(log n)

I binary-search for where the resulting elements start in the array. It’s the first index i so that arr[i] is better than arr[i+k] (with “better” meaning closer to or equally close to x). Then I just return the k elements starting there.

我们找到一个长度为k的序列，起始位置为i，如果arr[i] > arr[i+k]，说明从位置i+k可能存在可能的更优的解，此时l += 1；如果arr[i] <= arr[i+k] 说明i往左可能存在更优的解，此时r=mid。这里r的起始值为len(arr) - k，是为了确保m+k的时候不会超过arr的范围

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