第K个缺失的正整数代码出现TLE错误的原因及解决方法咨询

• 核心原因：flag变量未在每次外层循环中重置为0。只要某次循环中i存在于数组，flag被设为1后就会永久保持1，后续所有i都会被误判为存在于数组中，count无法继续累加，外层while(count < k)陷入死循环，最终触发超时。
• 次要效率问题：即使修复了flag重置的问题，当前算法时间复杂度也达到了O((k+n)logn)：外层最多循环k+n次（k和n最大都是1000，极限下要循环2000次），每次循环内部做一次O(logn)的二分查找，虽然本题约束下不会超时，但属于不必要的性能损耗。

1. 快速修复现有代码死循环问题 #

仅需要在每次外层循环开头把flag重置为0即可解决TLE，还可以做小优化减少不必要的运算：

class Solution {
public:
    int findKthPositive(vector<int>& arr, int k) {
        int count=0,i=1;
        while(count<k)
        {
            int flag=0; // 每次循环重置flag
            int start=0,end=arr.size()-1;
            while(start<=end)
            {
                int mid=start+(end-start)/2;
                if(arr[mid]==i) {
                    flag=1;
                    break; // 找到就提前退出二分，减少运算
                }
                else if(arr[mid]>i) 
                    end=mid-1; 
                else 
                    start=mid+1;
            }
            if(flag==0) count++;
            i++;
        } 
        return i-1; // 不用额外存储缺失数组，直接返回结果节省空间
    }
};

2. 更优的O(logn)算法实现 #

利用有序数组的性质可以大幅降低时间复杂度：到下标mid（从0开始）为止，缺失的正整数个数为arr[mid] - (mid + 1)（到mid位置共有mid+1个元素，正常无缺失的话最后一个元素应该是mid+1，实际值和理论值的差值就是缺失的正整数数量）。通过二分查找直接定位第k个缺失数的位置，代码如下：

class Solution {
public:
    int findKthPositive(vector<int>& arr, int k) {
        int left = 0, right = arr.size();
        while (left < right) {
            int mid = left + (right - left) / 2;
            if (arr[mid] - mid - 1 < k) {
                left = mid + 1;
            } else {
                right = mid;
            }
        }
        return left + k;
    }
};

内容的提问来源于stack exchange，提问作者Deepanshu Pathak