A hard binary search problem

658. Find K Closest Elements

We have four cases as shown in the comments. Each case, we will have different moving direction.

Pay attention, we are comparing

A[mid] to A[mid + k -1]

with A[mid + 1] to A[mid + k], this exaplains why we have arr[mid] and arr[mid + k] instead of arr[mid + k -1]

class Solution {
public:
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        int lo = 0, hi = arr.size() - k;
        /*
         1    x  mid         mid + k
         2       mid  x      mid + k
         3       mid       x mid + k
         4       mid         mid + k  x
         1: go left:  x - A[mid] <  A[mid + k] - x
         2: go left:  x - A[mid] <=  A[mid + k] - x
         3: go right: x - A[mid] > A[mid + k] - x
         4: go right: x - A[mid] > A[mid + k] - x
                
        */
        while (lo < hi) {
            int mid = (lo + hi) / 2;
            if (x - arr[mid] > arr[mid + k] - x)  lo = mid + 1;
            else hi = mid;
        }
        return vector<int> (arr.begin() + lo, arr.begin() + lo + k);  
    }
};

Or we can change it to a complicated one.

class Solution {
public:
     /*  FOUR cases, you can either go left or right
         1    x  [mid]         [mid + k]
         2       [mid]  x      [mid + k]
         3       [mid]       x [mid + k]
         4       [mid]         [mid + k]  x
         1: go left:  x <= arr[mid]
         2: go left:  arr[mid] <= x && x <= arr[mid + k] && x - arr[mid] <= arr[mid + k] - x
         3: go right: arr[mid] <= x && x <= arr[mid + k] && x - arr[mid] > arr[mid + k] - x
         4: go right: x >= arr[mid + k]      
        */
    // First check go Left
    vector<int> findClosestElementsGoLeft(vector<int>& arr, int k, int x) {
        int lo = 0, hi = arr.size() - k;
        while (lo < hi) 
        {
            int mid = (lo + hi) / 2;
            // case 1 and case 2, stay left
            if (x <= arr[mid] || (arr[mid] <= x && x <= arr[mid + k] && x - arr[mid] <= arr[mid + k] - x))
                hi = mid; 
            else lo = mid + 1;
        }
        return vector<int> (arr.begin() + lo, arr.begin() + lo + k);  
    }
    // First check go Right
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        int lo = 0, hi = arr.size() - k;
        while (lo < hi) 
        {
            int mid = (lo + hi) / 2;
            //if (x - arr[mid] > arr[mid + k] - x)  lo = mid + 1;
            if ((arr[mid] <= x && x <= arr[mid + k] && x - arr[mid] > arr[mid + k] - x) || x >= arr[mid + k]) lo = mid + 1;
            else hi = mid;
        }
        return vector<int> (arr.begin() + lo, arr.begin() + lo + k);  
    }
};

There is no syntax color for the above code. Check this one: