[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 104
3. Absolute value of elements in the array and x will not exceed 104

解題思路

題意爲從給定有序數組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);
    }
};