[LeetCode] 658. Find K Closest Elements

题目链接: https://leetcode.com/problems/find-k-closest-elements/description/

Description

Given a sorted array, 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: [1,2,3,4,5], k=4, x=3
Output: [1,2,3,4]

Example 2:

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

Note:
1. The value k is positive and will always be smaller than the length of the sorted array.
2. Length of the given array is positive and will not exceed $10^4$
3. Absolute value of elements in the array and x will not exceed $10^4$

解题思路

题意为从给定有序数组arr中找到与x最接近的k个数字。有序数组找指定数，首先想到的就是用二分查找来确定大致位置。二分查找算法通过修改一点点代码可以有很多不同种的预期值，我在这道题中使用的是，查找最后一个小于等于x的元素下标，从该下标开始向左向右搜索更接近该值的数，用两个变量left、right来表示边界下标，最终获得需要的区间为[left + 1, right)。

另外，由于二分查找是找最后一个小于等于x的数字，若x = -1，arr = [1, 2, 3]，即数组中没有比其更小的数字时，搜索得到的下标将为-1，需要调整一下边界值，令right = 1即可。

时间复杂度为 O(logn + k)，空间复杂度为 O(1)

Code

class Solution {
public:
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        int left = 0, right = arr.size() - 1;
        int middle;

        while (left <= right) {
            middle = (left + right) >> 1;
            if (x < arr[middle]) {
                right = middle - 1;
            } else {
                left = middle + 1;
            }
        }

        if (right < 0) right = 1;
        left = right - 1;

        while (k-- > 0) {
            if (left >= 0 && right < arr.size()) {
                if ((x - arr[left]) <= (arr[right] - x)) {
                    left--;
                } else {
                    right++;
                }
            } else if (left >= 0) {
                left--;
            } else {
                right++;
            }
        }    

        return vector<int>(arr.begin() + left + 1, arr.begin() + right);
    }
};