Leetcode 327. 区间和的个数 (前缀和 + 离散化 + 树状数组)

Leetcode 327. 区间和的个数 (前缀和 + 离散化 + 树状数组)

题目

题意

有多少个连续的子数组，其和在$[lower,upper]$之间

题解

可以想到的做法：用前缀和在$O(1)$查询$[i,j]$的和，枚举所有的二元组$[i,j]$, 满足条件就加上。

可以优化为：$Pre$为前缀和数组， 从小到大枚举 $j$， 由于$\text{lower}\le Pre[j]-Pre[i-1]\le \text{upper}$ ，可以得到 $P[i-1]$ 满足$Pre[j]-\text{upper}\le Pre[i-1]\le Pre[j]-\text{lower}$ ，通过枚举$j$，可以将$[Pre[j]-\text{upper},Pre[j]-\text{lower}]$ 看做$[L,R]$, 之后查询所有$[L,R]$内的个数即为答案。

• 前缀和

​ 使用前缀和算出子数组$[i,j]$的和$Pre[j]-Pre[i]$。

• 离散化

由于数据范围较大，因此可以通过离散化降低数据。我们可以将$Pre[j]-\text{upper},Pre[j]-\text{lower},Pre[j]$ 一起排序后进行离散化。

• 树状数组 / 线段树 / 平衡树

这些数据结构都满足在$O(logn)$ 的时间复杂度查询$[L,R]$内的和。

代码

#define ll long long
class Solution {
   
public:
    vector<int> tree;
    int n;
    int lowbits(int x){
   
        return x & (-x);    
    }
    void update(int x){
   
        while(x <= n){
   
            tree[x] += 1;
            x += lowbits(x);
        }
    }
    int query(int x){
   
        int res = 0;
        while(x){
   
            res += tree[x];
            x -= lowbits(x);
        }
        return res;
    }
    int countRangeSum(vector<int>& nums, int lower, int upper) {
   
        ll sums = 0;
        vector<ll> preSum = {
   0};
        for(int x : nums){
   
            sums += x;
            preSum.emplace_back(sums);
        }

        set<ll> st;
        for(auto x : preSum){
   
            st.insert(x - lower);
            st.insert(x);
            st.insert(x - upper);
        }

        // 离散化
        unordered_map<ll, int> p;
        int c = 0;
        for(auto x : st) p[x] = c++;

        int res = 0;
        n = p.size();
        tree = vector<int> (n+5, 0);
        // cout << n << endl;
        for(auto x : preSum){
   
            int left = p[x-upper], right = p[x-lower];
            res += query(right+1) - query(left);
            // cout << x << " " << right << " " << query(right+1) << " " << left << " " << query(left) << endl;
            update(p[x]+1);
        }
        return res;
    }
};