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

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

题目

题意

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

题解

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

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

• 前缀和

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

• 离散化

由于数据范围较大，因此可以通过离散化降低数据。我们可以将(Pre[j] - extit{upper}, Pre[j] - extit{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;
    }
};