Count of Range Sum

Given an integer array nums, return the number of range sums that lie in [lower, upper] inclusive.

Range sum S(i, j) is defined as the sum of the elements in nums between indices i and 
j (i ? j), inclusive.



Note:
A naive algorithm of O(n2) is trivial. You MUST do better than that.


Example:
Given nums = [-2, 5, -1], lower = -2, upper = 2,
Return 3.
The three ranges are : [0, 0], [2, 2], [0, 2] and their respective sums are: -2, -1, 2.

Credits:Special thanks to @dietpepsi for adding this problem and creating all test cases.

Solution

import java.util.Iterator;
import java.util.NavigableMap;
import java.util.TreeMap;
public class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        if (nums == null || nums.length == 0 || upper < lower) return 0;

        int size = nums.length;
        long[] cum = new long[size];
        cum[0] = nums[0];
        for (int i = 1; i < size; i++) {
            cum[i] = cum[i - 1] + nums[i];
        }

        TreeMap<Long, Integer> map = new TreeMap<>();
        map.put((long) 0,1);

        int result = 0;
        for (int i = 0; i < size; i++) {
            long lowerBound = cum[i] - upper;
            long higherBound = cum[i] - lower;

            NavigableMap<Long, Integer> navigableSet = map.subMap(lowerBound, true, higherBound, true);
            Iterator<Integer> iterator = navigableSet.values().iterator();
            while (iterator.hasNext()) {
                result += iterator.next();
            }

            map.merge(cum[i], 1, (a,b)->a+b);
        }

        return result;
    }

}