Contiguous subarray with sum in a range

Given an array int32 arr[] of size n, return the number of non-empty contigious subarrays whose sum lies in range [a, b]

Examples:

count([1,2,3], 0, 3) = 4 ( [1], [2], [3], [1, 2])
count([-2,5,-1], -2, 2) = 3 ( [-2], [-1], [-2, 5, -1] )

Naive O(n^2) solution –

def naive_algorithm(lst, a, b):
    result = 0
    for i in xrange(len(lst)):
        for j in xrange(i, len(lst)):
            if a <= sum(lst[i:j + 1]) <= b:
                result += 1
    return result

How we can improve? Think that if we were asked for contiguous subarray with sum = a then how would we would have solved it?
Yes, we can compute cumsum while traversing the array left to right and at each position j we would have checked if we have have a position i < j such that sum[i..j] = a. We can easily check this by checking the fact that sum[i..j] = cumsum[j]-cumsum[i-1] = a, i.e. cumsum[i-1]=cumsum[j]-a . That is we can cache the cumsum along the way and check at each j if (cumsum[j]-a) is already present in the cache. Now, we want to extend this idea for a rannge [a..b] instead of a.

A simple O(nlgn) solution would be using cumsum and Java navigable set’s subset operation. The idea is that if the number itself falls between the range than its a subset and we count it. Otherwise, if cumsum at the position of the element falls between the range then there is a subset from the start to this position and we count it.

Note that, there might be many other subarray’s in between start and current position that can sum to the given range. We can check if A[i..j] sums to S by checking if cumsum[j]-cumsum[i] = S. In other words, at each position i, we can check if cumsum[j]-S was seen earlier in some previous position i < j . We can easily do this using a hash to track each of the cumsum at each position.

Now, question is how to track the range of sum? Idea is to put the cumsums seen so far in a TreeSet (i.e. BST order) and for a position i we find how many elements fall in the range by finding number of elements in the subSet from (cumsum[i]-b) to (cumsum[i]-a) (think why?). So, we increase count by this number.

Below is a O(nlgn) implementation of this logic. Note that, treeset is handled in such a way that we don’t have to call size() method that is O(n) itself. Instead I keep index and calculate size by using simple formula of (endIndex-startIndex+1) .

Overall time complexity O(nlgn), overall space complexity O(n).

//find number of sub arrays with A<=sum<=rank 
public static int subArrayWithSumInRange(int[] A, int a, int b){
	int count = 0;
	TreeSet<Pair> treeSet = new TreeSet<test.Pair>();
	int cumsum = 0;
	
	for(int i = 0; i< A.length; i++){
		cumsum+=A[i];
		
		if(A[i] >= a && A[i] <= b){
			count++;
		}
		else if(cumsum >= a && cumsum <= b){
			count++;
		}

		NavigableSet<Pair> subSet = treeSet.subSet(new Pair(cumsum-b, i+1), true, new Pair(cumsum-a, i+1), false);
		if(!subSet.isEmpty()){
			count += Math.abs(subSet.first().second - subSet.last().second)+1;
		}
		treeSet.add(new Pair(cumsum, i));
	}
	
	return count;
}

private static class Pair implements Comparable<Pair>{
	public int first;
	public int second;
	
	public Pair(int first, int second){
		this.first = first;
		this.second = second;
	}

	@Override
	public int compareTo(Pair o) {
		if(this.first == o.first)
			return Integer.compare(this.second, o.second);
		else
			return Integer.compare(this.first, o.first);
	}

	@Override
	public String toString() {
		return "[" + first + "," + second + "]";
	}
}

Alternative Method – Merge Sort

Remember the problem we discussed in another post where we counted number of smaller elements on right of each element i.e. we counted the occurrences of x[j] < x[i], for j > i. That is we counted occurrences of x[j] – x[i] < 0, for j > i.

In this problem of range sum we need to count the occurrences of a<=(x[j]-x[i])<= b, for j > i

Therefore, we can solve this problem similar way we solved the count smaller on right problem, that is using merge sort in O(nlgn), where we count the answer while doing the merge. During the merge stage, we have already sorted the left half [start, mid) and right half [mid, end). We then iterate through the left half with index i. For each i, we need to find two indices sl and su in the right half where

sl is the first lower index that satisfies S[j] - S[i] > upper;
su is the first upper index that satisfies S[k] - S[i] >= lower.

Then the number of sums in [lower, upper] is (su-sl). We also use another index j and an auxiliary array merged[] to perform the actual merge sort. Below is the implementation of this idea that runs in O(nlgn) time and O(n) space.

public static int countRangeSum(int nums[], int lower, int upper){
	int[] cumsum = new int[nums.length];
	cumsum[0] = nums[0];
	for(int i = 1; i<nums.length; i++){
		cumsum[i] = cumsum[i-1]+nums[i];
	}
	
	return countRangeSumWithMerge(cumsum, 0, nums.length-1, lower, upper);
}

public static int countRangeSumWithMerge(int A[], int p, int r, int lower, int upper){
	int count = 0;
	int n = r-p+1;
	
	if(p >= r){
		return 0;
	}
	//divide (sort)
	int mid = (p+r)/2;
	count += countRangeSumWithMerge(A, p, mid, lower, upper);
	count += countRangeSumWithMerge(A, mid+1, r, lower, upper);
	
	//conquer (merge)
	int merged[] = new int[n];
	int sl = mid+1, su = mid+1; 
	for(int i = p, j = mid+1, k = 0; i <= mid && j<=r; k++){
		//count ranges
		while(sl <= r && (A[sl]-A[i]) < lower) sl++;
		while(su <= r && (A[su]-A[i]) <= upper) su++;
		count += su-sl;
		
		//merge sort
		if(A[i] < A[j]){
			merged[k] = A[i++];
		}
		else{
			merged[k] = A[j++];
		}
	}
	System.arraycopy(merged, 0, A, 0, merged.length);
	
	return count;
}