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;
    }
};