Leetcode 327. Count of Range Sum[hard]

进入hard模式。

题目：
Given an integer array nums, return the number of range sums that lie in [lower, upper] inclusive.
Range sum S(i, j) is defined as the sum of the elements in nums between indices i and j (i ≤ j), inclusive.

Note:
A naive algorithm of O(n2) is trivial. You MUST do better than that.

Example:
Given nums = [-2, 5, -1], lower = -2, upper = 2,
Return 3.
The three ranges are : [0, 0], [2, 2], [0, 2] and their respective sums are: -2, -1, 2.

---

区间求和的问题一般都会转化为前缀和。
设前i项的和为sum[i]，则我们需要找到前缀和sum[j]（0<=j< i），使得lower<=sum[i]-sum[j]<=upper。
即sum[i]-upper<=sum[j]<=sum[i]-lower。
也就是说我们需要在sum[0]~sum[i-1]中找有多少个符合以上条件。
很直观的想法就是平衡树：顺次向树中插入sum[i]，在插入sum[i]之前查找树中有多少个符合条件的sum[j]，将所有的这些答案求和就是答案。

但是我偷懒并没有写平衡树，而是使用树状数组取代平衡树的功能。
先将sum[i]离散化。之后顺次插入sum[i]，在每次插入sum[i]之前，在树状数组中查找小于等于sum[i]-lower的数字个数，和小于sum[i]-upper的数字个数，作差（和0取max）即为符合条件的sum[j]的个数。

class Solution {public:    int n;    vector<int> tr;    int countRangeSum(vector<int>& nums, int lower, int upper) {        long long sum = 0;        int ans = 0;        vector<long long> s, v;        s.push_back(0);        v.push_back(0);        for (int i = 0; i < nums.size(); i++) {            sum += nums[i];            s.push_back(sum);            v.push_back(sum);        }        sort(s.begin(), s.end());        n = v.size() - 1;        tr.clear();        tr.resize(n + 2);        for (int i = 1; i <= n; i++) {            int x = lower_bound(s.begin(), s.end(), v[i - 1]) - s.begin() + 1;            add(x);            int a = lower_bound(s.begin(), s.end(), v[i] - upper) - s.begin();            int b = upper_bound(s.begin(), s.end(), v[i] - lower) - s.begin();            if (a > b) continue;            ans += find(b) - find(a);        }        return ans;    }    void add(int x) {        while (x <= n + 1) {            tr[x]++;            x += x & -x;        }    }    int find(int x) {        int ans = 0;        while (x) {            ans += tr[x];            x -= x & -x;        }        return ans;    }};