LeetCode 327. Count of Range Sum

327. Count of Range Sum

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.

思路： 先算出sum[i], 然后找出sum[i]-sum[j]在upper和lower之间的数目。

hash+树状数组。 

AC代码：

class Solution {
public:
    int lobit(int x)
    {
        return x&(x^(x-1));
    }
    void add(int x, int pos, vector<int>& array)
    {
        while(pos<array.size())
        {
            array[pos] += x;
            pos += lobit(pos);
        }
    }
    int get(int pos, vector<int> & array)
    {
        int ret = 0;
        while(pos > 0)
        {
            ret += array[pos];
            pos -= lobit(pos);
        }
        return ret;
    }
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        int n = nums.size();
        if(n == 0) return 0;
        vector<long long> ele;
        ele.push_back(0);
        ele.push_back(-(1LL<<39));
        long long sum = 0;
        for(int i=0; i<n; i++)
        {
            sum += nums[i];
            ele.push_back( sum);
            ele.push_back( sum-lower);
            ele.push_back( sum-upper);
        }
        sort(ele.begin(), ele.end());
        map<long long,int> mp;
        for(int i=0; i<ele.size(); i++)
        {
            mp[ele[i]] = i+2;
        }
        vector<int>array(ele.size()+10, 0);
        add(1, mp[0], array);
        sum = 0;
        int ret = 0;
        for(int i=0; i<n; i++)
        {
            sum += nums[i];
            int left = mp[sum - upper];
            int right = mp[sum-lower];
            ret += get(right,array) - get(left-1, array);
            add(1, mp[sum], array);
        }
        return ret;
    }
};