Leetcode 327.区间和的个数

区间和的个数

给定一个整数数组 nums，返回区间和在 [lower, upper] 之间的个数，包含 lower 和 upper。
区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)。

说明:
最直观的算法复杂度是 O(n²) ，请在此基础上优化你的算法。

示例:

输入: nums = [-2,5,-1], lower = -2, upper = 2,

输出: 3

解释: 3个区间分别是: [0,0], [2,2], [0,2]，它们表示的和分别为: -2, -1, 2。

使用TreeMap来保存每个前缀和的计数

 import java.util.TreeMap;

 public class Solution {

     public static void main(String[] args){
         int[] nums={-2,5,-1};
         System.out.println(countRangeSum(nums,-2,2));
     }

     public static int countRangeSum(int[] nums, int lower, int upper) {
         if (nums == null || nums.length == 0) return 0;
         long[] sums = new long[nums.length];
         sums[0] = nums[0];
         for (int i = 1; i < nums.length; i++) sums[i] = sums[i - 1] + nums[i];
         int total = 0;
         TreeMap<Long, Integer> treemap = new TreeMap<>();
         for (int i = 0; i < nums.length; i++) {
             if (lower <= sums[i] && sums[i] <= upper) {
                 total++;
             }
             for (Integer count : treemap.subMap(sums[i] - upper, true, sums[i] - lower, true).values()) {
                 total += count;
             }
             Integer count = treemap.get(sums[i]);
             if (count == null) {
                 count = 1;
             }
             else {
                 count++;
             }
             treemap.put(sums[i], count);
         }
         return total;
     }
 }