【前缀和】A013_LC_区间和的个数（treemap+公式转换）

一、Problem

给定一个整数数组 nums，返回区间和在 [lower, upper] 之间的个数，包含 lower 和 upper

区间和 S(i, j) 表示在 nums 中，位置从 i 到 j 的元素之和，包含 i 和 j (i ≤ j)

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

输入: nums = [-2,5,-1], lower = -2, upper = 2,
输出: 3 
解释: 3个区间分别是: [0,0], [2,2], [0,2]，它们表示的和分别为: -2, -1, 2

二、Solution

方法一：公式优化

思考

最朴素的做法是双层循环在前缀和数组中找找 s[i]-s[j] ∈ [lo, hi] 的区间， $O(n^2)$，最常见的优化出发点是转换公式：lo ≤ s[i] - s[j] ≤ hi

$$\{\begin{array}{r}lo\le s[i]-s[j]\le hi①\\ lo-s[i]\le -s[j]\le hi-s[i]②\\ s[i]-hi\le s[j]\le s[i]-lo③\end{array}$$

所以查找符合③式要求的 s[j] 的个数就相当于代替了第二层循环在前缀和数组中找前缀和的作用

注意溢出

class Solution {
     
    public int countRangeSum(int[] A, int lo, int hi) {
     
        long n=A.length, s=0, ans=0;
        TreeMap<Long, Integer> mp = new TreeMap<>();
        mp.put(0l, 1);

        for (int x : A) {
     
            s+=x;
            for (int cnt : mp.subMap(s-hi, true, s-lo, true).values()) 
                ans+=cnt;
            mp.put(s, mp.getOrDefault(s, 0)+1);
        }
        return (int) ans;
    }
}

复杂度分析

• 时间复杂度：$O(nlogn)$，
• 空间复杂度：$O(n)$