LeetCode | 658. Find K Closest Elements

题目

Given a sorted integer array arr, two integers k and x, return the k closest integers to x in the array. The result should also be sorted in ascending order.

An integer a is closer to x than an integer b if:

• |a - x| < |b - x|, or
• |a - x| == |b - x| and a < b

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 <= 104
• arr is sorted in ascending order.
• -10^4 <= arr[i], x <= 10^4

代码

class Solution {
public:
    int findX(vector<int>& arr, int x) {
        int low = 0, high = arr.size() - 1;
        while(low < high)
        {
            int mid = (low + high) / 2;
            if(arr[mid] == x)
                return mid;
            if(arr[mid] < x)
                low = mid + 1;
            else
                high = mid;
        }
        if(arr[high] > x)
        {
            if(high > 0 && x - arr[high-1] < arr[high] - x)
                return high - 1;
        }
        return high;
    }
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        vector<int> res;
        if(arr.size() < 1)
            return res;
        if(x <= arr[0])
        {
            res = vector(arr.begin(), arr.begin()+k);
            return res;
        }
        if(x >= arr[arr.size() - 1])
        {
            res = vector(arr.end()-k, arr.end());
            return res;
        }
        int idx = findX(arr, x), startIdx = idx, endIdx = idx;
        for(int j = 1; j < k && startIdx >= 0 && endIdx < arr.size(); j++)
        {
            if(startIdx == 0)
            {
                endIdx++;
            }
            else if(endIdx == arr.size() - 1)
            {
                startIdx--;
            }
            else
            {
                if(x - arr[startIdx - 1] <= arr[endIdx + 1] - x)
                {
                    startIdx--;
                }
                else
                {
                    endIdx++;
                }
            }
        }
            
        //cout << idx << " " << startIdx << " " << endIdx << endl;
        res = vector(arr.begin()+startIdx, arr.begin()+endIdx+1);
        return res;
    }
};