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

最新推荐文章于 2022-09-13 16:06:03 发布

wdt_

最新推荐文章于 2022-09-13 16:06:03 发布

阅读量162

点赞数

文章标签： treemap

本文链接：https://blog.csdn.net/qq_43539599/article/details/107857562

版权

前缀和与差分专栏收录该内容

16 篇文章 0 订阅

订阅专栏

一、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

$\left\{ \begin{aligned} lo ≤ s[i] - s[j] ≤ hi\ ① \\ lo - s[i] ≤ -s[j] ≤ hi - s[i]\ ② \\ s[i] - hi ≤ s[j] ≤ s[i] - lo\ ③ \end{aligned} \right.$

所以查找符合③式要求的 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;
    }
}