Lintcode 404. Subarray Sum II

Given an integer array, find a subarray where the sum of numbers is in a given interval. Your code should return the number of possible answers. (The element in the array should be positive)

Example

Given [1,2,3,4] and interval = [1,3], return 4. The possible answers are:

[0, 0]
[0, 1]
[1, 1]
[2, 2]

Code (Java):

public class Solution {
    /**
     * @param A: An integer array
     * @param start: An integer
     * @param end: An integer
     * @return: the number of possible answer
     */
    public int subarraySumII(int[] A, int start, int end) {
        if (A == null || A.length == 0) {
            return 0;
        }
        
        int[] sum = new int[A.length + 1];
        for (int i = 1; i <= A.length; i++) {
            sum[i] += sum[i - 1] + A[i - 1];
        }
        
        int lo = 0;
        int hi = 0;
        
        int ans = 0;
        for (int j = 1; j <= A.length; j++) {
            while (lo < j && sum[lo] < sum[j] - end) {
                lo++;
            }
            
            while (hi < j && sum[hi] <= sum[j] - start) {
                hi++;
            }
            
            ans += hi - lo;
        }
        
        return ans;
    }
}