C++ solution using binary indexed tree


  • 0
    B
    class Solution {
    public:
    typedef long long LL;
    vector<LL> sum;
    vector<LL> bitree;
    
    int bit(int x) {
    	return x & -x;
    }
    
    void insert(int u, int val, int n) {
    	for (int i = u; i <= n; i += bit(i)) {
    		bitree[i] += val;
    	}
    }
    
    int query(int a, int b) {
    	return get(b) - get(a - 1);
    }
    
    int get(int x) {
    	int ret = 0;
    	for (int i = x; i > 0; i -= bit(i)) {
    		ret += bitree[i];
    	}
    	return ret;
    }
    
    int countRangeSum(vector<int>& nums, int lower, int upper) {
    	int n = nums.size();
    	sum.assign(3 * n + 1, 0LL);
    
    	for (int i = 1; i <= n; i++) {
    		sum[i] = sum[i - 1] + (LL)nums[i - 1];
    		sum[n + i] = (LL)lower + sum[i];
    		sum[2 * n + i] = (LL)upper + sum[i];
    	}
    
    	sum.push_back(lower);
    	sum.push_back(upper);
    
    	set<LL> aux(sum.begin(), sum.end());
    	vector<LL> vec(aux.begin(), aux.end());
    	map<LL, int> m;
    	for (int i = 0; i < vec.size(); i++) {
    		m[vec[i]] = i + 1;
    	}
    
    	int N = m.size();
    	bitree.assign(m.size() + 1, 0);
    	insert(m[sum[n]], 1, N);
    
    	int ret = 0;
    	for (int i = n; i >= 1; i--) {
    		ret += query(m[lower + sum[i - 1]], m[upper + sum[i - 1]]);
    		insert(m[sum[i - 1]], 1, N);
    	}
    
    	return ret;
    }
    

    };


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.