658. Find K Closest Elements

class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        int xIndex = searchNearestIndex(arr, x);
        int startIndex = xIndex, endIndex = xIndex; // [startIndex, endIndex]
        k--;
        while (k > 0) {
            if (startIndex - 1 < 0) {
                // 只能向右扩张
                endIndex++;
            } else if (endIndex + 1 >= arr.length) {
                // 只能向左扩张
                startIndex--;
            } else {
                if (x - arr[startIndex - 1] <= arr[endIndex + 1] - x) {
                    startIndex--;
                } else {
                    endIndex++;
                }
            }
            k--;
        }

        List<Integer> numList = new ArrayList<>(endIndex - startIndex + 1);
        for (int i = startIndex; i <= endIndex; i++) {
            numList.add(arr[i]);
        }
        return numList;
    }

    private int searchNearestIndex(int[] arr, int x) {
        // 首先通过二分搜索找到大于等于 x 的最小元素
        int i = 0, j = arr.length - 1;
        while (i < j) {
            int mid = (i + j) >>> 1;
            if (arr[mid] < x) {
                i = mid + 1;
            } else if (arr[mid] > x) {
                j = mid;
            } else {
                return mid;
            }
        }

        if (j == 0) {
            return 0;
        } else {
            // 再从左侧相邻元素与大于等于 x 的最小元素中选择最近的元素
            i = j - 1;
            if (Math.abs(arr[i] - x) <= Math.abs(arr[j] - x)) { // 注意此处是小于等于
                return i;
            } else {
                return j;
            }
        }
    }
}

References

658. Find K Closest Elements