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(n^2)$$ 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

Solution

class Solution {
public:
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        long sub_sum = 0;
        vector<long> sorted(1, 0);
        int res = 0;
        for ( int i = 0; i < int(nums.size()); i++ ) {
            sub_sum += nums[i];
            int lbd = lower_bound(sorted.begin(), sorted.end(), sub_sum - upper) - sorted.begin();
            int ubd = upper_bound(sorted.begin(), sorted.end(), sub_sum - lower) - sorted.begin() - 1;
            res += ubd - lbd + 1;
            int newPos = upper_bound(sorted.begin(), sorted.end(), sub_sum) - sorted.begin();
            sorted.insert(sorted.begin() + newPos, sub_sum);
        }
        return res;
    }
};

思路：
先把原数组转换成前缀数组和，然后进行扫描。扫描过程中利用二分法寻找range两个端点可以插入的位置，以及这个前缀数组和可插入的位置，相当于进行了一遍插入排序。（这里直接用了库函数upper_bound和lower_bound）。但可能insert的操作有$$O(N)$$复杂度，导致了整体time complexity的下降。（利用搜索二叉树来改进。）

class Solution {
public:
    struct TreeNode{
        long val;
        int nleft, nright, count;
        TreeNode *left, *right;
        TreeNode(long x): val(x), count(1), nleft(0), nright(0), left(NULL), right(NULL) {}
    };
    int bound(TreeNode* root, long x, bool tag) {
        // lbd: tag = flase; ubd: tag = true;
        TreeNode* p = root;
        int icount = 0;
        while ( p != NULL ) {
            if ( x == p->val ) return ( icount + p->nleft + tag*(p->count));
            if ( x < p->val ) p = p->left;
            else {
                icount += p->count + p->nleft;
                p = p->right;
            }
        }
        return icount;
    }
    void insert(TreeNode *root, long x) {
        TreeNode *p = root;
        while ( p != NULL ) {
            if ( x == p->val ) { p->count += 1; return; }
            if ( x < p->val ) {
                p->nleft += 1;
                if ( p->left == NULL ) { p->left = new TreeNode(x); return;}
                p = p->left;
            }
            else {
                p->nright += 1;
                if ( p->right == NULL ) { p->right = new TreeNode(x); return;}
                p = p->right;
            }
        }
        return;
    }
    int countRangeSum(vector<int>& nums, int lower, int upper) {
        int n = nums.size();
        if ( n == 0 ) return 0;
        long sub_sum = 0;
        TreeNode *root = new TreeNode(0);
        int res = 0;
        for ( int i = 0; i < n; i++ ) {
            sub_sum += nums[i];
            int lbd = bound(root, sub_sum - upper, false);
            int ubd = bound(root, sub_sum - lower, true);
            res += ubd - lbd;
            insert(root, sub_sum);
        }
        return res;
    }
};

• 注意upper_bound和lower_bound的用法

经过整理之后，这题的思路就非常清晰了。但看到很多利用“merge sort”的解法，感觉挺妙的，研究之。。。