327. Count of Range Sum (Hard)

https://leetcode.com/problems/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(n²) is trivial. You MUST do better than that.

Example:

Input: nums = [-2,5,-1], lower = -2, upper = 2,
Output: 3 
Explanation: The three ranges are : [0,0], [2,2], [0,2] and their respective sums are: -2, -1, 2.

Solutions

public class Solution {
    public int countRangeSum(int[] nums, int lower, int upper) {
        if (nums == null || nums.length == 0) {
            return 0;
        }

        int ans = 0;

        // sums of all preceding elements, use long in case overflow.
        long[] sums = new long[nums.length + 1];
        Arrays.fill(sums, 0);

        // keep the count of same sum, elements in TreeMap is sorted by key increasingly.
        TreeMap<Long, Integer> sumMap = new TreeMap<>();
        sumMap.put(0L, 1);

        for (int i = 0; i < nums.length; i++) {
            sums[i + 1] = sums[i] + nums[i];

            if (!sumMap.containsKey(sums[i + 1])) {
                sumMap.put(sums[i + 1], 0);
            }

            sumMap.put(sums[i + 1], sumMap.get(sums[i + 1]) + 1);
        }

        for (int i = 0; i < sums.length; i++) {

            // Find out all the ranges that start from i (inclusive) that are subject to lower<= rangeSum <=upper.

            // Given that, sums before i should be removed
            sumMap.put(sums[i], sumMap.get(sums[i]) - 1);
            if (sumMap.get(sums[i]) == 0) {
                sumMap.remove(sums[i]);
            }

            // range sum starting from i to x is sum[x] - sum[i], subject to lower < sum[x] - sum[i] < upper
            Map<Long, Integer> subMap = sumMap.subMap(sums[i] + lower, sums[i] + upper + 1);
            for (Long key: subMap.keySet()) {
                ans += subMap.get(key);
            }
        }

        return ans;
    }
}

Incorrect Solutions

References