[LeetCode 862] Shortest Subarray with Sum at Least K

Return the length of the shortest, non-empty, contiguous subarray of A with sum at least K.

If there is no non-empty subarray with sum at least K, return -1.

Example 1:

Input: A = [1], K = 1
Output: 1

Example 2:

Input: A = [1,2], K = 4
Output: -1

Example 3:

Input: A = [2,-1,2], K = 3
Output: 3

Note:

1 <= A.length <= 50000
-10 ^ 5 <= A[i] <= 10 ^ 5
1 <= K <= 10 ^ 9

Input size is up to 5 * 10^4, this makes sliding window to check all subarrays TLE as the runtime is O(N^2). To compute a subarray's sum, we use prefix sum techique S[L, R] = PS[R] - PS[L - 1]. The shortest subarray must end at one of the element in A, so we can go through A, check each shortest subarray length that ends at A[i]. To compute the shortest length of subarray that ends at A[i] and has total sum >= K, we need to find the largest index j such that PS[i] - PS[j] >= K. There are 2 key observations to help us derive an efficient solution.

1. for index j and i, j < i, if prefix sum PS[j] >= PS[i], then for any index k > i, we'll never consider [j, k] as a possible answer. This is because we can always cut off the [j, i - 1] subarray to get a NOT smaller sum and shorter length.

2. for index j and i, j < i, if [L1, j] is the shortest subarray with sum >= K and ends at A[j], then to get an even shorter answer that ends at A[i], the starting index L2 of such answer must be L2 > L1. Otherwise we'll have j - L1 < i - L2, this answer is worse than A[j]'s answer. This means the same prefix sum is used at most once to compute a candidate answer.

With the above analysis, we can come up with two efficient solutions.

Solution 1. O(N * logN) maintaining an increasing list + binary search solution.

Based on observation 1, if store a pair of sum and index for each A[i] in a list, we know that before we add a new pair, we always remove previous non-smaller sum. This list is always strictly increasing. So we can do a binary search to find the largest index j such that PS[j] <= currSum - K. Then remove previous non-smaller sum entries before appending the current entry to the end of this list.

class Solution {
    public int shortestSubarray(int[] A, int K) {
        long sum = 0;
        long best = A.length + 1;
        List<long[]> list = new ArrayList<>();
        list.add(new long[]{0, -1});
        
        for(int i = 0; i < A.length; i++) {
            sum += A[i];
            int j = binarySearch(list, sum - K);
            if(j >= 0) {
                best = Math.min(best, i - list.get(j)[1]);
            }
            while(list.size() > 0 && list.get(list.size() - 1)[0] >= sum) {
                list.remove(list.size() - 1);
            }
            list.add(new long[]{sum, i});
        }
        return best > A.length ? -1 : (int)best;
    }

    private int binarySearch(List<long[]> list, long target) {
        int l = 0, r = list.size() - 1;
        while(l < r - 1) {
            int mid = l + (r - l) / 2;
            long[] e = list.get(mid);
            if(e[0] > target) {
                r = mid - 1;
            }
            else {
                l = mid;
            }
        }
        if(list.get(r)[0] <= target) {
            return r;
        }
        else if(list.get(l)[0] <= target) {
            return l;
        }
        return -1;
    }
}

Solution 2. O(N) Deque solution.

An even better solution is to use both observations. We store all pairs in Deque and still remove last non-smaller sums before adding current new pair entry to the end of Deque. So the sums in this deque is still strictly increasing. However, instead of using binary search to find a possible answer, we can just poll from the head of Deque and check if the prefix sum difference is >= K. Keep polling from deque's head until the sum condition can not be met anymore. The last met condition index gives a possible answer for subarray ending at the current element. Observation point 2 tells us that once a prefix sum is used, it'll never be used again, so we do not need to add them back to the deque. Each possible prefix sum pair is added to deque once and polled out of deque at most once, thus the runtime is O(N).

class Solution {
    public int shortestSubarray(int[] A, int K) {       
        int n = A.length, best = A.length + 1;
        long sum = 0;
        ArrayDeque<Long> dq1 = new ArrayDeque<>();
        ArrayDeque<Integer> dq2 = new ArrayDeque<>();
        dq1.add(0L);
        dq2.add(-1);
        
        for(int i = 0; i < n; i++) {
            sum += A[i];
            int j = A.length;
            while(dq1.size() > 0 && dq1.peekFirst() <= sum - K) {
                dq1.pollFirst();
                j = dq2.pollFirst();                
            }
            if(j < A.length) {
                best = Math.min(best, i - j);
            }
            while(dq1.size() > 0 && dq1.peekLast() >= sum) {
                dq1.pollLast();
                dq2.pollLast();
            }
            dq1.addLast(sum);
            dq2.addLast(i);
        }
        return best > A.length ? -1 : best;
    }
}