Count of Range Sum

This page explains Java solution to problem Count of Range Sum using TreeMap data structure.

Problem Statement

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 1:

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.

Solution

If you have any suggestions in below code, please create a pull request by clicking here.

package com.vc.hard;

import java.util.*;

class CountOfRangeSum {
    public int countRangeSum(int[] nums, int lower, int upper) {
        TreeMap<Long, Integer> prefixSum = new TreeMap<>();
        long sum = 0;
        int count = 0;
        prefixSum.put(0l, 1);
        for(int i = 0; i < nums.length; i++) {
            sum += nums[i];
            long from = sum - upper;
            long to = sum - lower;
            Map<Long, Integer> subMap = prefixSum.subMap(from, true, to, true);
            for(int value: subMap.values()) count += value;
            prefixSum.put(sum, prefixSum.getOrDefault(sum, 0) + 1);
        }
        return count;
    }
}

Time Complexity

O(N * log(N)) Where
N is total number of elements in an input array

Space Complexity

O(N) Where
N is total number of elements in an input array