460.在排序数组中找最接近的K个数

描述

给一个目标数 target, 一个非负整数 k, 一个按照升序排列的数组 A。
在A中找与target最接近的k个整数。返回这k个数并按照与target的接近
程度从小到大排序，如果接近程度相当，那么小的数排在前面。

样例

如果 A = [1, 2, 3], target = 2 and k = 3, 那么返回 [2, 1, 3].
如果 A = [1, 4, 6, 8], target = 3 and k = 3, 那么返回 [4, 1, 6].

挑战

O(logn + k) 的时间复杂度

代码

public class Solution {
    /**
     * @param A an integer array
     * @param target an integer
     * @param k a non-negative integer
     * @return an integer array
     */
    public int[] kClosestNumbers(int[] A, int target, int k) {
        int[] result = new int[k];
        
        if (A == null || A.length == 0) {
            return A;
        }
        // 此处不能带=，因为等于时有输出顺序要求
        if (k > A.length) {
            return A;
        }
        
        int index = firstIndex(A, target);
        
        //  A[start] 和 A[end]是最接近target的两个数
        int start = index - 1;
        int end = index;
        for (int i = 0; i < k; i++) {
            // index = 0 时，start < 0，此处做异常检测
            if (start < 0) {
                result[i] = A[end++];
            // index = A.length - 1
            } else if (end >= A.length) {
                result[i] = A[start--];
            } else {
                // 接近程度相当的话，小的数在前面
                if (target - A[start] <= A[end] - target) {
                    result[i] = A[start--];
                } else {
                    result[i] = A[end++];
                }
            }
        }
        return result;
    }
    
    /* 寻找第一个位置，同样分三种情况：
     * target小于A[0]，target大于A[A.length - 1]，target在数组中
     */
    private int firstIndex(int[] A, int target) {
        int start = 0, end = A.length - 1;
        while (start + 1 < end) {
            int mid = start + (end - start) / 2;
            if (A[mid] < target) {
                start = mid;
            } else if (A[mid] > target) {
                end = mid;
            } else {
                end = mid;
            }
        }
        
        // A[start] == target 和 traget小于A[0]
        if (A[start] >= target) {
            return start;
        }
        // target在数组中，不包括A[start] == target情况
        if (A[end] >= target) {
            return end;
        }
        // target大于A[A.length - 1]
        return A.length;
    }
}