D - Kth Excluded Java

D - Kth Excluded Java

题目

https://atcoder.jp/contests/abc205/tasks/abc205_d

给N个正整数
会有Q次查询
每次查询给一个正整数k，需要返回第k小的正整数（并且不能是给的N个数中的数）

解决思路

核心思路：差分数列 + 二分法

1、计算差分数列

1 2 3 4 5 6 7 8 9 10 11 12 …

给定数列：A = [3, 6, 7, 9, 10]

差分数列：a = [2, 4, 4, 5, 5]

差分数列计算：ai = Ai - i (假设 i 从1开始，不同于代码中数组下标从0开始)

差分数列含义：ai表示Ai之前有多少不同于整个A数列的数 。

A1=3，3之前有1、2，一共有2个不同于A的数，因此a1=2；

A2=6，6之前有1、2、4、5，因此一共有4个不同于A的数，因此a2=4；

A3=7，7之前有1、2、4、5，因此一共有4个不同于A的数，因此a3=4

2、二分法找到k对应差分数列a的位置

寻找什么位置：需要找到a中小于k的元素中，最大的索引index，也就是说k是在A[index]到A[index + 1]之间的数，即a[index] < k <= a[index + 1]

k=3，则index=1，a[index] = a1 = 2；

k=4，则index=1，a[index] = a1 = 2；

k=5，则index=3，a[index] = a3 = 4；

计算含义：用一个具体的case举例，k=3时，需要寻找不同于A的第3小的一个数（可以直接看出来是4），那么我们根据上述规则寻找到index=1，在差分数列中对应a1=2，原数列A1=3，即A1之前只有2个可选的数，所以第3小一定在比A1大；index + 1 = 2，对应a2=4，A2=6，即A2之前有4个可选的数，所以第3小一定比A2小；这样就定位了k在A1和A2之间。

3、根据寻找到的位置计算结果

找到位置后如何计算结果：res = A[index] + (k - a[index])

k=3，则index=1，此时A1=3, a1=2，则res = 3 + （3 - 2） = 4；

k=4，则index=1，此时A1=3, a1=2，则res = 3 + （4 - 2） = 5；

k=5，则index=3，此时A3=7, a3=4，则res = 7 + （5 - 4） = 8；

计算含义：继续使用上述的case，k=3时，我们定位了k在A1和A2之间。A1之前有2个可选的数，那么k - ai = 3 - 2 = 1，这个1表明第k小的数比A1大1，因此res = A1 + 1 = 3 + 1 = 4；

疑问点：

不知道上述解决思路还有没有可以优化的地方？
下边附上代码，如何看自己是否满足要求（Time Limit: 3 sec / Memory Limit: 1024 MB）？代码在逻辑上、耗时上、内存上，还有没有可以优化的点呀？

Java代码

import com.alibaba.fastjson.JSON;

public class DKExcluded {

    /**
     * 二分法寻找某个值对应的index
     */
    public static int dichotomy(int[] arr, int k) {
        int start = 0;
        int end = arr.length - 1;
        while (start <= end) {
            int mid = (start + end) / 2;
            if (k == arr[mid]) {
                return mid;
            } else if (k < arr[mid]) {
                end = mid - 1;
            } else {
                start = mid + 1;
            }
        }
        return -1;
    }

    /**
     * 给N个正整数（不包含0）
     * 会有Q次查询
     * 每次查询给一个正整数k，需要返回第k小的正整数（并且不能是给的N个数中的数）
     */
    public static int[] findQKth(int[] arr, int[] qArr) {
        int[] kArr = new int[qArr.length];
        int i = 0;
        for (int q : qArr) {
            kArr[i] = findKth(arr, q);
            i++;
        }
        return kArr;
    }
    private static int findKth(int[] arr, int k) {
        int size = arr.length;;
        // 差分数列 diffArr[i]表示 arr[i]前有多少个不同于arr数列的数（不包含0）
        int[] diffArr = new int[size];
        for (int i = 0; i < size; i++) {
            diffArr[i] = arr[i] - (i + 1);
        }

        // 边界处理
        if (k <= diffArr[0]) {
            return k;
        }
        if (k > diffArr[size - 1]) {
            return arr[size - 1] + (k - diffArr[size - 1]);
        }

        // 二分法找到index
        int index = findIndex(diffArr, k);

        // 根据index计算第k小的正整数
        return calculateKth(arr, diffArr, k, index);
    }

    /**
     * 二分法
     * arr从小达到排序的数组，可能存在元素相同的情况
     * arr[0] < k <= arr[arr.length - 1]
     * 寻找小于k的元素中最大的index
     */
    private static int findIndex(int[] arr, int k) {
        int start = 0;
        int end = arr.length - 1;
        int mid;
        while (start <= end) {
            mid = (start + end) / 2;
            if (k <= arr[mid]) {
                if (k > arr[mid - 1]) {
                    return mid - 1;
                } else {
                    end = mid - 1;
                }
            } else { // k > arr[mid]
                start = mid + 1;
            }
        }
        return -1;
    }

    private static int calculateKth(int[] arr, int[] diffArr, int k, int index) {
        return arr[index] + k - diffArr[index];
    }

    public static void main(String[] args) {
        int[] arr = new int[] {3, 5, 6, 7};
        int[] qArr = new int[] {2, 5, 3};
        // 期望：2, 9, 4
        int[] kArr = findQKth(arr, qArr);
        System.out.println(JSON.toJSONString(kArr));

        int[] arr2 = new int[] {1, 2, 3, 4, 5};
        int[] qArr2 = new int[] {1, 10};
        // 期望：6, 15
        int[] kArr2 = findQKth(arr2, qArr2);
        System.out.println(JSON.toJSONString(kArr2));
    }