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

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

题目

题意

有多少个连续的子数组，其和在 [ l o w e r , u p p e r ] [lower, upper] [lower,upper]之间

题解

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

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

• 前缀和

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

• 离散化

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

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

这些数据结构都满足在 O ( l o g n ) O(logn) O(logn) 的时间复杂度查询 [ L , R ] [L, R] [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;
    }
};