327. 区间和的个数 - 《算法》

题目
思路
代码

题目

给你一个整数数组 nums 以及两个整数 lower 和 upper 。求数组中，值位于范围 [lower, upper] （包含 lower 和 upper）之内的区间和的个数。

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

示例 1：
输入：nums = [-2,5,-1], lower = -2, upper = 2
输出：3
解释：存在三个区间：[0,0]、[2,2] 和 [0,2] ，对应的区间和分别是：-2 、-1 、2 。

示例 2：

输入：nums = [0], lower = 0, upper = 0
输出：1

提示：

1 <= nums.length <= 10^5
-2^31 <= nums[i] <= 2^31 - 1
-10^5 <= lower <= upper <= 10^5
题目数据保证答案是一个 32 位的整数

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/count-of-range-sum
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

思路

再来练习一下树状数组。

求区间和首先考虑前缀和，在求出前缀和数组的情况下，还需要枚举子数组的始末下标，会超时，因此需要想一个nlogn的算法。

对于一个子数组[i,j]，其和等于preSum[j]-preSum[i]。对于下标j，题目需要求使得lower<=preSum[j]-preSum[i]<=upper成立的i的数目，将不等式变形一下，可以得到preSum[i]>=preSum[j]-upper和preSum[i]<=preSum[j]-lower，使这两个不等式同时成立的i的数目需要累计到答案中。

因此，问题的核心可以归结为：对于preSum[j]，求使preSum[i]落在[preSum[j]-upper, preSum[j]-lower]范围内的下标i的个数。

那么，可以枚举下标j，计算满足条件的i的数目，为了满足nlogn的时间复杂度，可以使用平衡树，或者树状数组。

在枚举下标j之前，需要将前缀和数组维护为有序离散的，有序是必要的，因为要查询指定范围内的元素个数，离散是因为节省空间（nums[i]太大，不节省也放不下）。

代码

class Solution {
    int len = 0;
    int[] tree;
    public int countRangeSum(int[] nums, int lower, int upper) {
        int n = nums.length;
        long[] pre  = new long[n + 1];
        for (int i = 1; i <= n; i++) {
            pre[i] = pre[i - 1] + nums[i - 1];
        }
        // 将前缀和数组排序去重，然后放在Map中离散化
        long[] arr = Arrays.stream(pre).distinct().sorted().toArray();
        // 因为下面用到了floorKey和ceilingKey，所以这里用了TreeMap
        TreeMap<Long, Integer> map = new TreeMap<>();
        for (int i = 0; i < arr.length; i++) {
            map.put(arr[i], i + 1);
        }
        len = map.size();
        tree = new int[len + 1];
        int ans = 0;
        for (int i = 0; i <= n; ++i) {
            // sum - lower和sum - upper可能不存在于map中
            // 对于上界就找一个不大于自身的最大的数，下界找一个不小于自身的最小的数，得到新的上下界[g,f]
            Long f = map.floorKey(pre[i] - lower);
            Long g = map.ceilingKey(pre[i] - upper);
            if (f != null && g != null) {
                // 求前缀和落在[g,f]内的个数，注意左边要减1,才会包括g处的
                ans += query(map.get(f)) - query(map.get(g) - 1);
            }
            // 前缀和为sum的个数+1
            // 注意，必须先更新ans再update，因为对于下标0来说，当前还没有遍历到第一个元素，没有子数组满足条件，如果先update就会使ans变多
            update(map.get(pre[i]), 1);
        }
        return ans;
    }
    public void update(int i, int delta) {
        while (i <= len) {
            tree[i] += delta;
            i += lowbit(i);
        }
    }
    public int query(int i) {
        int sum = 0;
        while (i > 0) {
            sum += tree[i];
            i -= lowbit(i);
        }
        return sum;
    }
    public static int lowbit(int x) {
        return x & (-x);
    }
}

也可以参考「官解」的做法，在map中加入pre[i] - upper和pre[i] - lower，这样就不用使用floorKey和ceilingKey了。

class Solution {
    int len = 0;
    int[] tree;
    public int countRangeSum(int[] nums, int lower, int upper) {
        int n = nums.length;
        long[] pre  = new long[n + 1];
        for (int i = 1; i <= n; i++) {
            pre[i] = pre[i - 1] + nums[i - 1];
        }
        Set<Long> set = new TreeSet<>();
        for (long num : pre) {
            set.add(num);
            set.add(num - upper);
            set.add(num - lower);
        }
        Map<Long, Integer> map = new HashMap<>();
        int cnt = 1;
        for (long num : set) {
            map.put(num, cnt++);
        }
        len = map.size();
        tree = new int[len + 1];
        int ans = 0;
        for (int i = 0; i <= n; ++i) {
            ans += query(map.get(pre[i] - lower)) - query(map.get(pre[i] - upper) - 1);
            // 前缀和为sum的个数+1
            // 注意，必须先更新ans再update，因为对于下标0来说，当前还没有遍历到第一个元素，没有子数组满足条件，如果先update就会使ans变多
            update(map.get(pre[i]), 1);
        }
        return ans;
    }
    public void update(int i, int delta) {
        while (i <= len) {
            tree[i] += delta;
            i += lowbit(i);
        }
    }
    public int query(int i) {
        int sum = 0;
        while (i > 0) {
            sum += tree[i];
            i -= lowbit(i);
        }
        return sum;
    }
    public static int lowbit(int x) {
        return x & (-x);
    }
}