每日一题：区间和的个数

题目

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

Solution

维护前缀和，对于每个前缀和 $sum[i]$，等价于 $sum[0…i-1]$ 里面有多少个数满足差值在 $[lower, upper]$ 之间。

$sum[i] - sum[j] = [lower, upper]$ 等价于求出现在区间 $sum[j] \in [sum[i] - upper, sum[i] - lower]$ 的数的个数，故可以用树状数组维护。

遍历前缀和数组，对于当前前缀和 $sum[i]$，满足的个数便为：$query(sum[i] - lower) - query(sum[i] - upper - 1)$，累加答案即可。

由于数很大，故需要将所有出现的数离散化。

Code

class BIT {
    public:
        int* c;
        int n;
    public:
        BIT(int _n) {
            n = _n;
            c = new int[n + 1];
            for (int i = 0; i <= n; i++) c[i] = 0;
        }
        static constexpr int lowbit(int x) {
            return x & (-x);
        }
        void add(int x) {
            while (x <= n) {
                c[x]++;
                x += lowbit(x);
            }
        }
        int query(int x) {
            int sum = 0;
            while (x) {
                sum += c[x];
                x -= lowbit(x);
            }
            return sum;
        }
};

class Solution {
public:
    int countRangeSum(vector<int>& a, int lower, int upper) {
        set<long long> s;
        int m = a.size();
        long long *sum = new long long[m + 1];
        sum[0] = 0;
        for (int i = 0; i <= m; i++) {
            if (i) sum[i] = sum[i - 1] + a[i - 1];
            s.insert(sum[i]);
            s.insert(sum[i] - upper);
            s.insert(sum[i] - lower);
        }
        unordered_map<long long, int> mp;
        int n = s.size(), idx = 0;
        BIT bit(n);
        for (auto x : s) mp[x] = ++idx;
        int res = 0;
        for (int i = 0; i <= m; i++) {
            int l = mp[sum[i] - upper], r = mp[sum[i] - lower];
            res += bit.query(r) - bit.query(l - 1);
            bit.add(mp[sum[i]]);
        }
        return res;
    }
};